Add ref cleanup function

packages/react-dom/src/__tests__/refs-test.js

@@ -351,35 +351,6 @@ describe('ref swapping', () => {
       'Expected ref to be a function, a string, an object returned by React.createRef(), or null.',
     );
   });
-
-  it('should warn about callback refs returning a function', () => {
-    const container = document.createElement('div');
-    expect(() => {
-      ReactDOM.render(<div ref={() => () => {}} />, container);
-    }).toErrorDev('Unexpected return value from a callback ref in div');
-
-    // Cleanup should warn, too.
-    expect(() => {
-      ReactDOM.render(<span />, container);
-    }).toErrorDev('Unexpected return value from a callback ref in div', {
-      withoutStack: true,
-    });
-
-    // No warning when returning non-functions.
-    ReactDOM.render(<p ref={() => ({})} />, container);
-    ReactDOM.render(<p ref={() => null} />, container);
-    ReactDOM.render(<p ref={() => undefined} />, container);
-
-    // Still warns on functions (not deduped).
-    expect(() => {
-      ReactDOM.render(<div ref={() => () => {}} />, container);
-    }).toErrorDev('Unexpected return value from a callback ref in div');
-    expect(() => {
-      ReactDOM.unmountComponentAtNode(container);
-    }).toErrorDev('Unexpected return value from a callback ref in div', {
-      withoutStack: true,
-    });
-  });
 });
 
 describe('root level refs', () => {
@@ -534,3 +505,136 @@ describe('strings refs across renderers', () => {
     expect(inst.refs.child2).toBe(div2.firstChild);
   });
 });
+
+describe('refs return clean up function', () => {
+  it('calls clean up function if it exists', () => {
+    const container = document.createElement('div');
+    let cleanUp = jest.fn();
+    let setup = jest.fn();
+
+    ReactDOM.render(
+      <div
+        ref={_ref => {
+          setup(_ref);
+          return cleanUp;
+        }}
+      />,
+      container,
+    );
+
+    ReactDOM.render(
+      <div
+        ref={_ref => {
+          setup(_ref);
+        }}
+      />,
+      container,
+    );
+
+    expect(setup).toHaveBeenCalledTimes(2);
+    expect(cleanUp).toHaveBeenCalledTimes(1);
+    expect(cleanUp.mock.calls[0][0]).toBe(undefined);
+
+    ReactDOM.render(<div ref={_ref => {}} />, container);
+
+    expect(cleanUp).toHaveBeenCalledTimes(1);
+    expect(setup).toHaveBeenCalledTimes(3);
+    expect(setup.mock.calls[2][0]).toBe(null);
+
+    cleanUp = jest.fn();
+    setup = jest.fn();
+
+    ReactDOM.render(
+      <div
+        ref={_ref => {
+          setup(_ref);
+          return cleanUp;
+        }}
+      />,
+      container,
+    );
+
+    expect(setup).toHaveBeenCalledTimes(1);
+    expect(cleanUp).toHaveBeenCalledTimes(0);
+
+    ReactDOM.render(
+      <div
+        ref={_ref => {
+          setup(_ref);
+          return cleanUp;
+        }}
+      />,
+      container,
+    );
+
+    expect(setup).toHaveBeenCalledTimes(2);
+    expect(cleanUp).toHaveBeenCalledTimes(1);
+  });
+
+  it('handles ref functions with stable identity', () => {
+    const container = document.createElement('div');
+    const cleanUp = jest.fn();
+    const setup = jest.fn();
+
+    function _onRefChange(_ref) {
+      setup(_ref);
+      return cleanUp;
+    }
+
+    ReactDOM.render(<div ref={_onRefChange} />, container);
+
+    expect(setup).toHaveBeenCalledTimes(1);
+    expect(cleanUp).toHaveBeenCalledTimes(0);
+
+    ReactDOM.render(
+      <div className="niceClassName" ref={_onRefChange} />,
+      container,
+    );
+
+    expect(setup).toHaveBeenCalledTimes(1);
+    expect(cleanUp).toHaveBeenCalledTimes(0);
+
+    ReactDOM.render(<div />, container);
+
+    expect(setup).toHaveBeenCalledTimes(1);
+    expect(cleanUp).toHaveBeenCalledTimes(1);
+  });
+
+  it('warns if clean up function is returned when called with null', () => {
+    const container = document.createElement('div');
+    const cleanUp = jest.fn();
+    const setup = jest.fn();
+    let returnCleanUp = false;
+
+    ReactDOM.render(
+      <div
+        ref={_ref => {
+          setup(_ref);
+          if (returnCleanUp) {
+            return cleanUp;
+          }
+        }}
+      />,
+      container,
+    );
+
+    expect(setup).toHaveBeenCalledTimes(1);
+    expect(cleanUp).toHaveBeenCalledTimes(0);
+
+    returnCleanUp = true;
+
+    expect(() => {
+      ReactDOM.render(
+        <div
+          ref={_ref => {
+            setup(_ref);
+            if (returnCleanUp) {
+              return cleanUp;
+            }
+          }}
+        />,
+        container,
+      );
+    }).toErrorDev('Unexpected return value from a callback ref in div');
+  });
+});

packages/react-reconciler/src/ReactFiber.new.js

@@ -148,6 +148,7 @@ function FiberNode(
   this.index = 0;
 
   this.ref = null;
+  this.refCleanup = null;
 
   this.pendingProps = pendingProps;
   this.memoizedProps = null;
@@ -336,6 +337,7 @@ export function createWorkInProgress(current: Fiber, pendingProps: any): Fiber {
   workInProgress.sibling = current.sibling;
   workInProgress.index = current.index;
   workInProgress.ref = current.ref;
+  workInProgress.refCleanup = current.refCleanup;
 
   if (enableProfilerTimer) {
     workInProgress.selfBaseDuration = current.selfBaseDuration;
@@ -880,6 +882,7 @@ export function assignFiberPropertiesInDEV(
   target.sibling = source.sibling;
   target.index = source.index;
   target.ref = source.ref;
+  target.refCleanup = source.refCleanup;
   target.pendingProps = source.pendingProps;
   target.memoizedProps = source.memoizedProps;
   target.updateQueue = source.updateQueue;

packages/react-reconciler/src/ReactFiber.old.js

@@ -148,6 +148,7 @@ function FiberNode(
   this.index = 0;
 
   this.ref = null;
+  this.refCleanup = null;
 
   this.pendingProps = pendingProps;
   this.memoizedProps = null;
@@ -336,6 +337,7 @@ export function createWorkInProgress(current: Fiber, pendingProps: any): Fiber {
   workInProgress.sibling = current.sibling;
   workInProgress.index = current.index;
   workInProgress.ref = current.ref;
+  workInProgress.refCleanup = current.refCleanup;
 
   if (enableProfilerTimer) {
     workInProgress.selfBaseDuration = current.selfBaseDuration;
@@ -880,6 +882,7 @@ export function assignFiberPropertiesInDEV(
   target.sibling = source.sibling;
   target.index = source.index;
   target.ref = source.ref;
+  target.refCleanup = source.refCleanup;
   target.pendingProps = source.pendingProps;
   target.memoizedProps = source.memoizedProps;
   target.updateQueue = source.updateQueue;

packages/react-reconciler/src/ReactFiberCommitWork.new.js

@@ -286,8 +286,32 @@ function safelyAttachRef(current: Fiber, nearestMountedAncestor: Fiber | null) {
 
 function safelyDetachRef(current: Fiber, nearestMountedAncestor: Fiber | null) {
   const ref = current.ref;
+  const refCleanup = current.refCleanup;
+
   if (ref !== null) {
-    if (typeof ref === 'function') {
+    if (typeof refCleanup === 'function') {
+      try {
+        if (shouldProfile(current)) {
+          try {
+            startLayoutEffectTimer();
+            refCleanup();
+          } finally {
+            recordLayoutEffectDuration(current);
+          }
+        } else {
+          refCleanup();
+        }
+      } catch (error) {
+        captureCommitPhaseError(current, nearestMountedAncestor, error);
+      } finally {
+        // `refCleanup` has been called. Nullify all references to it to prevent double invocation.
+        current.refCleanup = null;
+        const finishedWork = current.alternate;
+        if (finishedWork != null) {
+          finishedWork.refCleanup = null;
+        }
+      }
+    } else if (typeof ref === 'function') {
       let retVal;
       try {
         if (shouldProfile(current)) {
@@ -1572,25 +1596,15 @@ function commitAttachRef(finishedWork: Fiber) {
       instanceToUse = instance;
     }
     if (typeof ref === 'function') {
-      let retVal;
       if (shouldProfile(finishedWork)) {
         try {
           startLayoutEffectTimer();
-          retVal = ref(instanceToUse);
+          finishedWork.refCleanup = ref(instanceToUse);
         } finally {
           recordLayoutEffectDuration(finishedWork);
         }
       } else {
-        retVal = ref(instanceToUse);
-      }
-      if (__DEV__) {
-        if (typeof retVal === 'function') {
-          console.error(
-            'Unexpected return value from a callback ref in %s. ' +
-              'A callback ref should not return a function.',
-            getComponentNameFromFiber(finishedWork),
-          );
-        }
+        finishedWork.refCleanup = ref(instanceToUse);
       }
     } else {
       if (__DEV__) {
@@ -1609,27 +1623,6 @@ function commitAttachRef(finishedWork: Fiber) {
   }
 }
 
-function commitDetachRef(current: Fiber) {
-  const currentRef = current.ref;
-  if (currentRef !== null) {
-    if (typeof currentRef === 'function') {
-      if (shouldProfile(current)) {
-        try {
-          startLayoutEffectTimer();
-          currentRef(null);
-        } finally {
-          recordLayoutEffectDuration(current);
-        }
-      } else {
-        currentRef(null);
-      }
-    } else {
-      // $FlowFixMe unable to narrow type to the non-function case
-      currentRef.current = null;
-    }
-  }
-}
-
 function detachFiberMutation(fiber: Fiber) {
   // Cut off the return pointer to disconnect it from the tree.
   // This enables us to detect and warn against state updates on an unmounted component.
@@ -4450,7 +4443,6 @@ function invokePassiveEffectUnmountInDEV(fiber: Fiber): void {
 export {
   commitPlacement,
   commitAttachRef,
-  commitDetachRef,
   invokeLayoutEffectMountInDEV,
   invokeLayoutEffectUnmountInDEV,
   invokePassiveEffectMountInDEV,