[Synchronous Suspense] Reuse deletions from primary tree

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');
+    });
   });
 });