[Synchronous Suspense] Reuse deletions from primary tree (#14133)

Fixes a bug where deletion effects in the primary tree were dropped
before entering the second render pass.

Because we no longer reset the effect list after the first render pass,
I've also moved the deletion of the fallback children to the complete
phase, after the tree successfully renders without suspending.

Will need to revisit this heuristic when we implement resuming.

Author: acdlite

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -1127,11 +1127,7 @@ function updateSuspenseComponent(
           progressedState !== null
             ? (workInProgress.child: any).child
             : (workInProgress.child: any);
-        reuseProgressedPrimaryChild(
-          workInProgress,
-          primaryChildFragment,
-          progressedPrimaryChild,
-        );
+        primaryChildFragment.child = progressedPrimaryChild;
       }
 
       const fallbackChildFragment = createFiberFromFragment(
@@ -1186,11 +1182,7 @@ function updateSuspenseComponent(
               ? (workInProgress.child: any).child
               : (workInProgress.child: any);
           if (progressedPrimaryChild !== currentPrimaryChildFragment.child) {
-            reuseProgressedPrimaryChild(
-              workInProgress,
-              primaryChildFragment,
-              progressedPrimaryChild,
-            );
+            primaryChildFragment.child = progressedPrimaryChild;
           }
         }
 
@@ -1213,20 +1205,19 @@ function updateSuspenseComponent(
         // and remove the intermediate fragment fiber.
         const nextPrimaryChildren = nextProps.children;
         const currentPrimaryChild = currentPrimaryChildFragment.child;
-        const currentFallbackChild = currentFallbackChildFragment.child;
         const primaryChild = reconcileChildFibers(
           workInProgress,
           currentPrimaryChild,
           nextPrimaryChildren,
           renderExpirationTime,
         );
-        // Delete the fallback children.
-        reconcileChildFibers(
-          workInProgress,
-          currentFallbackChild,
-          null,
-          renderExpirationTime,
-        );
+
+        // If this render doesn't suspend, we need to delete the fallback
+        // children. Wait until the complete phase, after we've confirmed the
+        // fallback is no longer needed.
+        // TODO: Would it be better to store the fallback fragment on
+        // the stateNode?
+
         // Continue rendering the children, like we normally do.
         child = next = primaryChild;
       }
@@ -1258,11 +1249,7 @@ function updateSuspenseComponent(
             progressedState !== null
               ? (workInProgress.child: any).child
               : (workInProgress.child: any);
-          reuseProgressedPrimaryChild(
-            workInProgress,
-            primaryChildFragment,
-            progressedPrimaryChild,
-          );
+          primaryChildFragment.child = progressedPrimaryChild;
         }
 
         // Create a fragment from the fallback children, too.
@@ -1298,49 +1285,6 @@ function updateSuspenseComponent(
   return next;
 }
 
-function reuseProgressedPrimaryChild(
-  workInProgress: Fiber,
-  primaryChildFragment: Fiber,
-  progressedChild: Fiber | null,
-) {
-  // This function is only called outside concurrent mode. Usually, if a work-
-  // in-progress primary tree suspends, we throw it out and revert back to
-  // current. Outside concurrent mode, though, we commit the suspended work-in-
-  // progress, even though it didn't complete. This function reuses the children
-  // and transfers the effects.
-  let child = (primaryChildFragment.child = progressedChild);
-  while (child !== null) {
-    // Ensure that the first and last effect of the parent corresponds
-    // to the children's first and last effect.
-    if (primaryChildFragment.firstEffect === null) {
-      primaryChildFragment.firstEffect = child.firstEffect;
-    }
-    if (child.lastEffect !== null) {
-      if (primaryChildFragment.lastEffect !== null) {
-        primaryChildFragment.lastEffect.nextEffect = child.firstEffect;
-      }
-      primaryChildFragment.lastEffect = child.lastEffect;
-    }
-
-    // Append all the effects of the subtree and this fiber onto the effect
-    // list of the parent. The completion order of the children affects the
-    // side-effect order.
-    if (child.effectTag > PerformedWork) {
-      if (primaryChildFragment.lastEffect !== null) {
-        primaryChildFragment.lastEffect.nextEffect = child;
-      } else {
-        primaryChildFragment.firstEffect = child;
-      }
-      primaryChildFragment.lastEffect = child;
-    }
-    child.return = primaryChildFragment;
-    child = child.sibling;
-  }
-
-  workInProgress.firstEffect = primaryChildFragment.firstEffect;
-  workInProgress.lastEffect = primaryChildFragment.lastEffect;
-}
-
 function updatePortalComponent(
   current: Fiber | null,
   workInProgress: Fiber,

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -82,6 +82,7 @@ import {
   popHydrationState,
 } from './ReactFiberHydrationContext';
 import {ConcurrentMode, NoContext} from './ReactTypeOfMode';
+import {reconcileChildFibers} from './ReactChildFiber';
 
 function markUpdate(workInProgress: Fiber) {
   // Tag the fiber with an update effect. This turns a Placement into
@@ -700,12 +701,29 @@ function completeWork(
       if ((workInProgress.effectTag & DidCapture) !== NoEffect) {
         // Something suspended. Re-render with the fallback children.
         workInProgress.expirationTime = renderExpirationTime;
-        workInProgress.firstEffect = workInProgress.lastEffect = null;
+        // Do not reset the effect list.
         return workInProgress;
       }
 
       const nextDidTimeout = nextState !== null;
       const prevDidTimeout = current !== null && current.memoizedState !== null;
+
+      if (current !== null && !nextDidTimeout && prevDidTimeout) {
+        // We just switched from the fallback to the normal children. Delete
+        // the fallback.
+        // TODO: Would it be better to store the fallback fragment on
+        // the stateNode during the begin phase?
+        const currentFallbackChild: Fiber | null = (current.child: any).sibling;
+        if (currentFallbackChild !== null) {
+          reconcileChildFibers(
+            workInProgress,
+            currentFallbackChild,
+            null,
+            renderExpirationTime,
+          );
+        }
+      }
+
       // The children either timed out after previously being visible, or
       // were restored after previously being hidden. Schedule an effect
       // to update their visiblity.

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -993,7 +993,6 @@ function completeUnitOfWork(workInProgress: Fiber): Fiber | null {
 
       if (nextUnitOfWork !== null) {
         // Completing this fiber spawned new work. Work on that next.
-        nextUnitOfWork.firstEffect = nextUnitOfWork.lastEffect = null;
         return nextUnitOfWork;
       }
 

packages/react-reconciler/src/__tests__/ReactSuspense-test.internal.js

@@ -779,5 +779,64 @@ describe('ReactSuspense', () => {
       expect(root).toFlushAndYield(['Child 1', 'Child 2']);
       expect(root).toMatchRenderedOutput(['Child 1', 'Child 2'].join(''));
     });
+
+    it('reuses effects, including deletions, from the suspended tree', () => {
+      const {useState} = React;
+
+      let setTab;
+      function App() {
+        const [tab, _setTab] = useState(0);
+        setTab = _setTab;
+
+        return (
+          <Suspense fallback={<Text text="Loading..." />}>
+            <AsyncText key={tab} text={'Tab: ' + tab} ms={1000} />
+            <Text key={tab + 'sibling'} text=" + sibling" />
+          </Suspense>
+        );
+      }
+
+      const root = ReactTestRenderer.create(<App />);
+      expect(ReactTestRenderer).toHaveYielded([
+        'Suspend! [Tab: 0]',
+        ' + sibling',
+        'Loading...',
+      ]);
+      expect(root).toMatchRenderedOutput('Loading...');
+      jest.advanceTimersByTime(1000);
+      expect(ReactTestRenderer).toHaveYielded([
+        'Promise resolved [Tab: 0]',
+        'Tab: 0',
+      ]);
+      expect(root).toMatchRenderedOutput('Tab: 0 + sibling');
+
+      setTab(1);
+      expect(ReactTestRenderer).toHaveYielded([
+        'Suspend! [Tab: 1]',
+        ' + sibling',
+        'Loading...',
+      ]);
+      expect(root).toMatchRenderedOutput('Loading...');
+      jest.advanceTimersByTime(1000);
+      expect(ReactTestRenderer).toHaveYielded([
+        'Promise resolved [Tab: 1]',
+        'Tab: 1',
+      ]);
+      expect(root).toMatchRenderedOutput('Tab: 1 + sibling');
+
+      setTab(2);
+      expect(ReactTestRenderer).toHaveYielded([
+        'Suspend! [Tab: 2]',
+        ' + sibling',
+        'Loading...',
+      ]);
+      expect(root).toMatchRenderedOutput('Loading...');
+      jest.advanceTimersByTime(1000);
+      expect(ReactTestRenderer).toHaveYielded([
+        'Promise resolved [Tab: 2]',
+        'Tab: 2',
+      ]);
+      expect(root).toMatchRenderedOutput('Tab: 2 + sibling');
+    });
   });
 });