Fix: Don't read primaryChild.childExpirationTime (#18457)

This is a variant of the fix in 5a0f1d. We can't rely on the primary
fiber's `childExpirationTime` field to be correct.

In this case, we can read from the Suspense boundary fiber instead.
This will include updates that exist in the fallback fiber, but that's
not a big deal; the important thing is that we don't drop updates.

Author: acdlite

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -1646,16 +1646,9 @@ function shouldRemainOnFallback(
 function getRemainingWorkInPrimaryTree(
   current: Fiber,
   workInProgress: Fiber,
-  currentPrimaryChildFragment: Fiber | null,
   renderExpirationTime,
 ) {
-  const currentParentOfPrimaryChildren =
-    currentPrimaryChildFragment !== null
-      ? currentPrimaryChildFragment
-      : current;
-  const currentChildExpirationTime =
-    currentParentOfPrimaryChildren.childExpirationTime;
-
+  const currentChildExpirationTime = current.childExpirationTime;
   const currentSuspenseState: SuspenseState = current.memoizedState;
   if (currentSuspenseState !== null) {
     // This boundary already timed out. Check if this render includes the level
@@ -1945,7 +1938,6 @@ function updateSuspenseComponent(
             primaryChildFragment.childExpirationTime = getRemainingWorkInPrimaryTree(
               current,
               workInProgress,
-              null,
               renderExpirationTime,
             );
             workInProgress.memoizedState = updateSuspenseState(
@@ -2016,7 +2008,6 @@ function updateSuspenseComponent(
         primaryChildFragment.childExpirationTime = getRemainingWorkInPrimaryTree(
           current,
           workInProgress,
-          currentPrimaryChildFragment,
           renderExpirationTime,
         );
         // Skip the primary children, and continue working on the
@@ -2118,7 +2109,6 @@ function updateSuspenseComponent(
         primaryChildFragment.childExpirationTime = getRemainingWorkInPrimaryTree(
           current,
           workInProgress,
-          null,
           renderExpirationTime,
         );
         // Skip the primary children, and continue working on the

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

@@ -3266,6 +3266,90 @@ describe('ReactSuspenseWithNoopRenderer', () => {
     },
   );
 
+  it(
+    'regression: primary fragment fiber is not always part of setState ' +
+      'return path (another case)',
+    async () => {
+      // Reproduces a bug where updates inside a suspended tree are dropped
+      // because the fragment fiber we insert to wrap the hidden children is not
+      // part of the return path, so it doesn't get marked during setState.
+      const {useState} = React;
+      const root = ReactNoop.createRoot();
+
+      function Parent() {
+        return (
+          <Suspense fallback={<Text text="Loading..." />}>
+            <Child />
+          </Suspense>
+        );
+      }
+
+      let setText;
+      function Child() {
+        const [text, _setText] = useState('A');
+        setText = _setText;
+        return <AsyncText text={text} />;
+      }
+
+      // Mount an initial tree. Resolve A so that it doesn't suspend.
+      await resolveText('A');
+      await ReactNoop.act(async () => {
+        root.render(<Parent />);
+      });
+      expect(Scheduler).toHaveYielded(['A']);
+      // At this point, the setState return path follows current fiber.
+      expect(root).toMatchRenderedOutput(<span prop="A" />);
+
+      // Schedule another update. This will "flip" the alternate pairs.
+      await resolveText('B');
+      await ReactNoop.act(async () => {
+        setText('B');
+      });
+      expect(Scheduler).toHaveYielded(['B']);
+      // Now the setState return path follows the *alternate* fiber.
+      expect(root).toMatchRenderedOutput(<span prop="B" />);
+
+      // Schedule another update. This time, we'll suspend.
+      await ReactNoop.act(async () => {
+        setText('C');
+      });
+      expect(Scheduler).toHaveYielded(['Suspend! [C]', 'Loading...']);
+
+      // Commit. This will insert a fragment fiber to wrap around the component
+      // that triggered the update.
+      await ReactNoop.act(async () => {
+        await advanceTimers(250);
+      });
+      // The fragment fiber is part of the current tree, but the setState return
+      // path still follows the alternate path. That means the fragment fiber is
+      // not part of the return path.
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span hidden={true} prop="B" />
+          <span prop="Loading..." />
+        </>,
+      );
+
+      await ReactNoop.act(async () => {
+        // Schedule a normal pri update. This will suspend again.
+        setText('D');
+
+        // And another update at lower priority. This will unblock.
+        await resolveText('E');
+        Scheduler.unstable_runWithPriority(
+          Scheduler.unstable_IdlePriority,
+          () => {
+            setText('E');
+          },
+        );
+      });
+      // Even though the fragment fiber is not part of the return path, we should
+      // be able to finish rendering.
+      expect(Scheduler).toHaveYielded(['Suspend! [D]', 'E']);
+      expect(root).toMatchRenderedOutput(<span prop="E" />);
+    },
+  );
+
   it(
     'after showing fallback, should not flip back to primary content until ' +
       'the update that suspended finishes',