Add ref cleanup function

Author: sammy-SC

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

@@ -534,3 +534,69 @@ describe('strings refs across renderers', () => {
     expect(inst.refs.child2).toBe(div2.firstChild);
   });
 });
+
+describe('refs return clean up function', () => {
+  it('does not break', () => {
+    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);
+  });
+});

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 = null;
 
   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 = null;
 
   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,7 +286,29 @@ function safelyAttachRef(current: Fiber, nearestMountedAncestor: Fiber | null) {
 
 function safelyDetachRef(current: Fiber, nearestMountedAncestor: Fiber | null) {
   const ref = current.ref;
-  if (ref !== null) {
+  const refCleanup = current.refCleanup;
+
+  if (refCleanup !== null) {
+    if (typeof ref === 'function') {
+      try {
+        if (shouldProfile(current)) {
+          try {
+            startLayoutEffectTimer();
+            refCleanup();
+          } finally {
+            recordLayoutEffectDuration(current);
+          }
+        } else {
+          refCleanup();
+        }
+      } catch (error) {
+        captureCommitPhaseError(current, nearestMountedAncestor, error);
+      }
+    } else {
+      // $FlowFixMe unable to narrow type to RefObject
+      ref.refCleanup = null;
+    }
+  } else if (ref !== null) {
     if (typeof ref === 'function') {
       let retVal;
       try {
@@ -1583,14 +1605,9 @@ function commitAttachRef(finishedWork: Fiber) {
       } 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),
-          );
-        }
+
+      if (typeof retVal === 'function') {
+        finishedWork.refCleanup = retVal;
       }
     } else {
       if (__DEV__) {

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

@@ -286,7 +286,29 @@ function safelyAttachRef(current: Fiber, nearestMountedAncestor: Fiber | null) {
 
 function safelyDetachRef(current: Fiber, nearestMountedAncestor: Fiber | null) {
   const ref = current.ref;
-  if (ref !== null) {
+  const refCleanup = current.refCleanup;
+
+  if (refCleanup !== null) {
+    if (typeof ref === 'function') {
+      try {
+        if (shouldProfile(current)) {
+          try {
+            startLayoutEffectTimer();
+            refCleanup();
+          } finally {
+            recordLayoutEffectDuration(current);
+          }
+        } else {
+          refCleanup();
+        }
+      } catch (error) {
+        captureCommitPhaseError(current, nearestMountedAncestor, error);
+      }
+    } else {
+      // $FlowFixMe unable to narrow type to RefObject
+      ref.refCleanup = null;
+    }
+  } else if (ref !== null) {
     if (typeof ref === 'function') {
       let retVal;
       try {
@@ -1583,14 +1605,9 @@ function commitAttachRef(finishedWork: Fiber) {
       } 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),
-          );
-        }
+
+      if (typeof retVal === 'function') {
+        finishedWork.refCleanup = retVal;
       }
     } else {
       if (__DEV__) {

packages/react-reconciler/src/ReactInternalTypes.js

@@ -131,6 +131,8 @@ export type Fiber = {
     | (((handle: mixed) => void) & {_stringRef: ?string, ...})
     | RefObject,
 
+  refCleanup: null | (() => void),
+
   // Input is the data coming into process this fiber. Arguments. Props.
   pendingProps: any, // This type will be more specific once we overload the tag.
   memoizedProps: any, // The props used to create the output.