Make sure Offscreen's ref is detached when unmounted

Author: sammy-SC

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

@@ -1121,7 +1121,7 @@ function commitLayoutEffectOnFiber(
           if (flags & Ref) {
             safelyAttachRef(finishedWork, finishedWork.return);
           }
-        } else if (finishedWork.pendingProps.mode !== undefined) {
+        } else {
           safelyDetachRef(finishedWork, finishedWork.return);
         }
       } else {
@@ -2148,6 +2148,8 @@ function commitDeletionEffectsOnFiber(
         offscreenSubtreeWasHidden =
           prevOffscreenSubtreeWasHidden || deletedFiber.memoizedState !== null;
 
+        safelyDetachRef(deletedFiber, nearestMountedAncestor);
+
         recursivelyTraverseDeletionEffects(
           finishedRoot,
           nearestMountedAncestor,
@@ -2833,6 +2835,7 @@ export function disappearLayoutEffects(finishedWork: Fiber) {
         // Nested Offscreen tree is already hidden. Don't disappear
         // its effects.
       } else {
+        safelyDetachRef(finishedWork, finishedWork.return);
         recursivelyTraverseDisappearLayoutEffects(finishedWork);
       }
       break;
@@ -2974,6 +2977,8 @@ export function reappearLayoutEffects(
           finishedWork,
           includeWorkInProgressEffects,
         );
+
+        safelyAttachRef(finishedWork, finishedWork.return);
       }
       break;
     }

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

@@ -1121,7 +1121,7 @@ function commitLayoutEffectOnFiber(
           if (flags & Ref) {
             safelyAttachRef(finishedWork, finishedWork.return);
           }
-        } else if (finishedWork.pendingProps.mode !== undefined) {
+        } else {
           safelyDetachRef(finishedWork, finishedWork.return);
         }
       } else {
@@ -2148,6 +2148,8 @@ function commitDeletionEffectsOnFiber(
         offscreenSubtreeWasHidden =
           prevOffscreenSubtreeWasHidden || deletedFiber.memoizedState !== null;
 
+        safelyDetachRef(deletedFiber, nearestMountedAncestor);
+
         recursivelyTraverseDeletionEffects(
           finishedRoot,
           nearestMountedAncestor,
@@ -2833,6 +2835,7 @@ export function disappearLayoutEffects(finishedWork: Fiber) {
         // Nested Offscreen tree is already hidden. Don't disappear
         // its effects.
       } else {
+        safelyDetachRef(finishedWork, finishedWork.return);
         recursivelyTraverseDisappearLayoutEffects(finishedWork);
       }
       break;
@@ -2974,6 +2977,8 @@ export function reappearLayoutEffects(
           finishedWork,
           includeWorkInProgressEffects,
         );
+
+        safelyAttachRef(finishedWork, finishedWork.return);
       }
       break;
     }

packages/react-reconciler/src/__tests__/ReactOffscreen-test.js

@@ -1307,4 +1307,86 @@ describe('ReactOffscreen', () => {
       expect(offscreenRef.current).not.toBeNull();
     });
   });
+
+  // @gate enableOffscreen
+  it('should detach ref if Offscreen is unmounted', async () => {
+    let offscreenRef;
+
+    function App({showOffscreen}) {
+      offscreenRef = useRef(null);
+      return showOffscreen ? (
+        <Offscreen
+          mode={'manual'}
+          ref={ref => {
+            offscreenRef.current = ref;
+          }}>
+          <div />
+        </Offscreen>
+      ) : null;
+    }
+
+    const root = ReactNoop.createRoot();
+
+    await act(async () => {
+      root.render(<App showOffscreen={true} />);
+    });
+
+    expect(offscreenRef.current).not.toBeNull();
+
+    await act(async () => {
+      root.render(<App showOffscreen={false} />);
+    });
+
+    expect(offscreenRef.current).toBeNull();
+
+    await act(async () => {
+      root.render(<App showOffscreen={true} />);
+    });
+
+    expect(offscreenRef.current).not.toBeNull();
+  });
+
+  // @gate enableOffscreen
+  it('should detach ref if Offscreen is unmounted', async () => {
+    let offscreenRef;
+    let divRef;
+
+    function App({mode}) {
+      offscreenRef = useRef(null);
+      divRef = useRef(null);
+      return (
+        <Offscreen mode={mode}>
+          <Offscreen
+            mode={'manual'}
+            ref={offscreenRef}>
+            <div ref={divRef} />
+          </Offscreen>
+        </Offscreen>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+
+    await act(async () => {
+      root.render(<App mode={'hidden'} />);
+    });
+
+    expect(offscreenRef.current).toBeNull();
+    expect(divRef.current).toBeNull();
+
+    await act(async () => {
+      root.render(<App mode={'visible'} />);
+    });
+
+    expect(offscreenRef.current).not.toBeNull();
+    expect(divRef.current).not.toBeNull();
+
+    console.log('Becoming Hidden');
+    await act(async () => {
+      root.render(<App mode={'hidden'} />);
+    });
+
+    expect(offscreenRef.current).toBeNull();
+    expect(divRef.current).toBeNull();
+  });
 });