Bugfix: memo drops lower pri updates on bail out (#18091)

Fixes a bug where lower priority updates on a components wrapped with
`memo` are sometimes left dangling in the queue without ever being
processed, if they are preceded by a higher priority bailout.

Cause
-----

The pending update priority field is cleared at the beginning of
`beginWork`. If there is remaining work at a lower priority level, it's
expected that it will be accumulated on the work-in-progress fiber
during the begin phase.

There's an exception where this assumption doesn't hold:
SimpleMemoComponent contains a bailout that occurs *before* the
component is evaluated and the update queues are processed, which means
we don't accumulate the next priority level. When we complete the fiber,
the work loop is left to believe that there's no remaining work.

Mitigation
----------

Since this only happens in a single case, a late bailout in
SimpleMemoComponent, I've mitigated the bug in that code path by
restoring the original update priority from the current fiber.

This same case does not apply to MemoComponent, because MemoComponent
fibers do not contain hooks or update queues; rather, they wrap around
an inner fiber that may contain those. However, I've added a test case
for MemoComponent to protect against a possible future regression.

Possible next steps
-------------------

We should consider moving the update priority assignment in `beginWork`
out of the common path and into each branch, to avoid similar bugs in
the future.

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -523,6 +523,20 @@ function updateSimpleMemoComponent(
     ) {
       didReceiveUpdate = false;
       if (updateExpirationTime < renderExpirationTime) {
+        // The pending update priority was cleared at the beginning of
+        // beginWork. We're about to bail out, but there might be additional
+        // updates at a lower priority. Usually, the priority level of the
+        // remaining updates is accumlated during the evaluation of the
+        // component (i.e. when processing the update queue). But since since
+        // we're bailing out early *without* evaluating the component, we need
+        // to account for it here, too. Reset to the value of the current fiber.
+        // NOTE: This only applies to SimpleMemoComponent, not MemoComponent,
+        // because a MemoComponent fiber does not have hooks or an update queue;
+        // rather, it wraps around an inner component, which may or may not
+        // contains hooks.
+        // TODO: Move the reset at in beginWork out of the common path so that
+        // this is no longer necessary.
+        workInProgress.expirationTime = current.expirationTime;
         return bailoutOnAlreadyFinishedWork(
           current,
           workInProgress,
@@ -3103,7 +3117,11 @@ function beginWork(
     didReceiveUpdate = false;
   }
 
-  // Before entering the begin phase, clear the expiration time.
+  // Before entering the begin phase, clear pending update priority.
+  // TODO: This assumes that we're about to evaluate the component and process
+  // the update queue. However, there's an exception: SimpleMemoComponent
+  // sometimes bails out later in the begin phase. This indicates that we should
+  // move this assignment out of the common path and into each branch.
   workInProgress.expirationTime = NoWork;
 
   switch (workInProgress.tag) {

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

@@ -419,6 +419,75 @@ describe('memo', () => {
           'Invalid prop `inner` of type `boolean` supplied to `Inner`, expected `number`.',
         ]);
       });
+
+      it('does not drop lower priority state updates when bailing out at higher pri (simple)', async () => {
+        const {useState} = React;
+
+        let setCounter;
+        const Counter = memo(() => {
+          const [counter, _setCounter] = useState(0);
+          setCounter = _setCounter;
+          return counter;
+        });
+
+        function App() {
+          return (
+            <Suspense fallback="Loading...">
+              <Counter />
+            </Suspense>
+          );
+        }
+
+        const root = ReactNoop.createRoot();
+        await ReactNoop.act(async () => {
+          root.render(<App />);
+        });
+        expect(root).toMatchRenderedOutput('0');
+
+        await ReactNoop.act(async () => {
+          setCounter(1);
+          ReactNoop.discreteUpdates(() => {
+            root.render(<App />);
+          });
+        });
+        expect(root).toMatchRenderedOutput('1');
+      });
+
+      it('does not drop lower priority state updates when bailing out at higher pri (complex)', async () => {
+        const {useState} = React;
+
+        let setCounter;
+        const Counter = memo(
+          () => {
+            const [counter, _setCounter] = useState(0);
+            setCounter = _setCounter;
+            return counter;
+          },
+          (a, b) => a.complexProp.val === b.complexProp.val,
+        );
+
+        function App() {
+          return (
+            <Suspense fallback="Loading...">
+              <Counter complexProp={{val: 1}} />
+            </Suspense>
+          );
+        }
+
+        const root = ReactNoop.createRoot();
+        await ReactNoop.act(async () => {
+          root.render(<App />);
+        });
+        expect(root).toMatchRenderedOutput('0');
+
+        await ReactNoop.act(async () => {
+          setCounter(1);
+          ReactNoop.discreteUpdates(() => {
+            root.render(<App />);
+          });
+        });
+        expect(root).toMatchRenderedOutput('1');
+      });
     });
   }
 });