[SuspenseList] Store lastEffect before rendering (#17131)

* Add a failing test for SuspenseList bug

* Store lastEffect before rendering

We can't reset the effect list to null because we don't rereconcile the
children so we drop deletion effects if we do that.

Instead we store the last effect as it was before we started rendering
so we can go back to where it was when we reset it.

We actually already do something like this when we delete the last row
for the tail="hidden" mode so we had a field available for it already.

Author: sebmarkbage

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -2355,18 +2355,20 @@ function initSuspenseListRenderState(
   tail: null | Fiber,
   lastContentRow: null | Fiber,
   tailMode: SuspenseListTailMode,
+  lastEffectBeforeRendering: null | Fiber,
 ): void {
   let renderState: null | SuspenseListRenderState =
     workInProgress.memoizedState;
   if (renderState === null) {
-    workInProgress.memoizedState = {
+    workInProgress.memoizedState = ({
       isBackwards: isBackwards,
       rendering: null,
       last: lastContentRow,
       tail: tail,
       tailExpiration: 0,
       tailMode: tailMode,
-    };
+      lastEffect: lastEffectBeforeRendering,
+    }: SuspenseListRenderState);
   } else {
     // We can reuse the existing object from previous renders.
     renderState.isBackwards = isBackwards;
@@ -2375,6 +2377,7 @@ function initSuspenseListRenderState(
     renderState.tail = tail;
     renderState.tailExpiration = 0;
     renderState.tailMode = tailMode;
+    renderState.lastEffect = lastEffectBeforeRendering;
   }
 }
 
@@ -2456,6 +2459,7 @@ function updateSuspenseListComponent(
           tail,
           lastContentRow,
           tailMode,
+          workInProgress.lastEffect,
         );
         break;
       }
@@ -2487,6 +2491,7 @@ function updateSuspenseListComponent(
           tail,
           null, // last
           tailMode,
+          workInProgress.lastEffect,
         );
         break;
       }
@@ -2497,6 +2502,7 @@ function updateSuspenseListComponent(
           null, // tail
           null, // last
           undefined,
+          workInProgress.lastEffect,
         );
         break;
       }

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -1053,7 +1053,10 @@ function completeWork(
                 // Rerender the whole list, but this time, we'll force fallbacks
                 // to stay in place.
                 // Reset the effect list before doing the second pass since that's now invalid.
-                workInProgress.firstEffect = workInProgress.lastEffect = null;
+                if (renderState.lastEffect === null) {
+                  workInProgress.firstEffect = null;
+                }
+                workInProgress.lastEffect = renderState.lastEffect;
                 // Reset the child fibers to their original state.
                 resetChildFibers(workInProgress, renderExpirationTime);
 

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

@@ -585,6 +585,90 @@ describe('ReactSuspenseList', () => {
     );
   });
 
+  it('displays all "together" during an update', async () => {
+    let A = createAsyncText('A');
+    let B = createAsyncText('B');
+    let C = createAsyncText('C');
+    let D = createAsyncText('D');
+
+    function Foo({step}) {
+      return (
+        <SuspenseList revealOrder="together">
+          {step === 0 && (
+            <Suspense fallback={<Text text="Loading A" />}>
+              <A />
+            </Suspense>
+          )}
+          {step === 0 && (
+            <Suspense fallback={<Text text="Loading B" />}>
+              <B />
+            </Suspense>
+          )}
+          {step === 1 && (
+            <Suspense fallback={<Text text="Loading C" />}>
+              <C />
+            </Suspense>
+          )}
+          {step === 1 && (
+            <Suspense fallback={<Text text="Loading D" />}>
+              <D />
+            </Suspense>
+          )}
+        </SuspenseList>
+      );
+    }
+
+    // Mount
+    await A.resolve();
+    ReactNoop.render(<Foo step={0} />);
+    expect(Scheduler).toFlushAndYield([
+      'A',
+      'Suspend! [B]',
+      'Loading B',
+      'Loading A',
+      'Loading B',
+    ]);
+    expect(ReactNoop).toMatchRenderedOutput(
+      <>
+        <span>Loading A</span>
+        <span>Loading B</span>
+      </>,
+    );
+    await B.resolve();
+    expect(Scheduler).toFlushAndYield(['A', 'B']);
+    expect(ReactNoop).toMatchRenderedOutput(
+      <>
+        <span>A</span>
+        <span>B</span>
+      </>,
+    );
+
+    // Update
+    await C.resolve();
+    ReactNoop.render(<Foo step={1} />);
+    expect(Scheduler).toFlushAndYield([
+      'C',
+      'Suspend! [D]',
+      'Loading D',
+      'Loading C',
+      'Loading D',
+    ]);
+    expect(ReactNoop).toMatchRenderedOutput(
+      <>
+        <span>Loading C</span>
+        <span>Loading D</span>
+      </>,
+    );
+    await D.resolve();
+    expect(Scheduler).toFlushAndYield(['C', 'D']);
+    expect(ReactNoop).toMatchRenderedOutput(
+      <>
+        <span>C</span>
+        <span>D</span>
+      </>,
+    );
+  });
+
   it('avoided boundaries can be coordinate with SuspenseList', async () => {
     let A = createAsyncText('A');
     let B = createAsyncText('B');