From 1608017e77274b378a421db9c071295fbe65fbae Mon Sep 17 00:00:00 2001
From: Samuel Susla <samuel.susla@gmail.com>
Date: Wed, 21 Sep 2022 11:18:53 +0100
Subject: [PATCH] Make sure Offscreen's ref is detached when unmounted

---
 .../src/ReactFiberCommitWork.new.js           |  7 +-
 .../src/ReactFiberCommitWork.old.js           |  7 +-
 .../src/__tests__/ReactOffscreen-test.js      | 80 +++++++++++++++++++
 3 files changed, 92 insertions(+), 2 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiberCommitWork.new.js b/packages/react-reconciler/src/ReactFiberCommitWork.new.js
index 5f8de4fab94e5..dcf07571ded4d 100644
--- a/packages/react-reconciler/src/ReactFiberCommitWork.new.js
+++ b/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;
     }
diff --git a/packages/react-reconciler/src/ReactFiberCommitWork.old.js b/packages/react-reconciler/src/ReactFiberCommitWork.old.js
index c3c2e964ad497..0fb539a5b4785 100644
--- a/packages/react-reconciler/src/ReactFiberCommitWork.old.js
+++ b/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;
     }
diff --git a/packages/react-reconciler/src/__tests__/ReactOffscreen-test.js b/packages/react-reconciler/src/__tests__/ReactOffscreen-test.js
index 6180c8a1e7105..c3dcb38e5ec52 100644
--- a/packages/react-reconciler/src/__tests__/ReactOffscreen-test.js
+++ b/packages/react-reconciler/src/__tests__/ReactOffscreen-test.js
@@ -1307,4 +1307,84 @@ 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();
+  });
 });