[Bugfix] Updates "un-committed" when rebasing

packages/react-reconciler/src/ReactFiberHooks.js

@@ -20,6 +20,7 @@ import type {SuspenseConfig} from './ReactFiberSuspenseConfig';
 import type {ReactPriorityLevel} from './SchedulerWithReactIntegration';
 
 import ReactSharedInternals from 'shared/ReactSharedInternals';
+import {rebasedUpdatesNeverUncommit} from 'shared/ReactFeatureFlags';
 
 import {NoWork} from './ReactFiberExpirationTime';
 import {readContext} from './ReactFiberNewContext';
@@ -144,8 +145,12 @@ if (__DEV__) {
 export type Hook = {
   memoizedState: any,
 
+  // TODO: These fields are only used by the state and reducer hooks. Consider
+  // moving them to a separate object.
   baseState: any,
   baseUpdate: Update<any, any> | null,
+  rebaseEnd: Update<any, any> | null,
+  rebaseTime: ExpirationTime,
   queue: UpdateQueue<any, any> | null,
 
   next: Hook | null,
@@ -580,6 +585,8 @@ function mountWorkInProgressHook(): Hook {
     baseState: null,
     queue: null,
     baseUpdate: null,
+    rebaseEnd: null,
+    rebaseTime: NoWork,
 
     next: null,
   };
@@ -621,6 +628,8 @@ function updateWorkInProgressHook(): Hook {
       baseState: currentHook.baseState,
       queue: currentHook.queue,
       baseUpdate: currentHook.baseUpdate,
+      rebaseEnd: currentHook.rebaseEnd,
+      rebaseTime: currentHook.rebaseTime,
 
       next: null,
     };
@@ -735,8 +744,10 @@ function updateReducer<S, I, A>(
   // The last update in the entire queue
   const last = queue.last;
   // The last update that is part of the base state.
-  const baseUpdate = hook.baseUpdate;
   const baseState = hook.baseState;
+  const baseUpdate = hook.baseUpdate;
+  const rebaseEnd = hook.rebaseEnd;
+  const rebaseTime = hook.rebaseTime;
 
   // Find the first unprocessed update.
   let first;
@@ -755,19 +766,36 @@ function updateReducer<S, I, A>(
     let newState = baseState;
     let newBaseState = null;
     let newBaseUpdate = null;
+    let newRebaseTime = NoWork;
+    let newRebaseEnd = null;
     let prevUpdate = baseUpdate;
     let update = first;
-    let didSkip = false;
+
+    // Track whether the update is part of a rebase.
+    // TODO: Should probably split this into two separate loops, instead of
+    // using a boolean.
+    let isRebasing = rebaseTime !== NoWork;
+
     do {
       const updateExpirationTime = update.expirationTime;
-      if (updateExpirationTime < renderExpirationTime) {
+      if (prevUpdate === rebaseEnd) {
+        isRebasing = false;
+      }
+      if (
+        // Check if this update should be skipped
+        updateExpirationTime < renderExpirationTime &&
+        // If we're currently rebasing, don't skip this update if we already
+        // committed it.
+        (!rebasedUpdatesNeverUncommit ||
+          (!isRebasing || updateExpirationTime < rebaseTime))
+      ) {
         // Priority is insufficient. Skip this update. If this is the first
         // skipped update, the previous update/state is the new base
         // update/state.
-        if (!didSkip) {
-          didSkip = true;
+        if (newRebaseTime === NoWork) {
           newBaseUpdate = prevUpdate;
           newBaseState = newState;
+          newRebaseTime = renderExpirationTime;
         }
         // Update the remaining priority in the queue.
         if (updateExpirationTime > remainingExpirationTime) {
@@ -787,7 +815,6 @@ function updateReducer<S, I, A>(
           updateExpirationTime,
           update.suspenseConfig,
         );
-
         // Process this update.
         if (update.eagerReducer === reducer) {
           // If this update was processed eagerly, and its reducer matches the
@@ -802,9 +829,11 @@ function updateReducer<S, I, A>(
       update = update.next;
     } while (update !== null && update !== first);
 
-    if (!didSkip) {
-      newBaseUpdate = prevUpdate;
+    if (newRebaseTime === NoWork) {
       newBaseState = newState;
+      newBaseUpdate = prevUpdate;
+    } else {
+      newRebaseEnd = prevUpdate;
     }
 
     // Mark that the fiber performed work, but only if the new state is
@@ -814,8 +843,10 @@ function updateReducer<S, I, A>(
     }
 
     hook.memoizedState = newState;
-    hook.baseUpdate = newBaseUpdate;
     hook.baseState = newBaseState;
+    hook.baseUpdate = newBaseUpdate;
+    hook.rebaseEnd = newRebaseEnd;
+    hook.rebaseTime = newRebaseTime;
 
     queue.lastRenderedState = newState;
   }

packages/react-reconciler/src/ReactUpdateQueue.js

@@ -97,7 +97,10 @@ import {
 import {Callback, ShouldCapture, DidCapture} from 'shared/ReactSideEffectTags';
 import {ClassComponent} from 'shared/ReactWorkTags';
 
-import {debugRenderPhaseSideEffectsForStrictMode} from 'shared/ReactFeatureFlags';
+import {
+  debugRenderPhaseSideEffectsForStrictMode,
+  rebasedUpdatesNeverUncommit,
+} from 'shared/ReactFeatureFlags';
 
 import {StrictMode} from './ReactTypeOfMode';
 import {
@@ -138,6 +141,9 @@ export type UpdateQueue<State> = {
 
   firstCapturedEffect: Update<State> | null,
   lastCapturedEffect: Update<State> | null,
+
+  rebaseTime: ExpirationTime,
+  rebaseEnd: Update<State> | null,
 };
 
 export const UpdateState = 0;
@@ -172,6 +178,8 @@ export function createUpdateQueue<State>(baseState: State): UpdateQueue<State> {
     lastEffect: null,
     firstCapturedEffect: null,
     lastCapturedEffect: null,
+    rebaseTime: NoWork,
+    rebaseEnd: null,
   };
   return queue;
 }
@@ -194,6 +202,9 @@ function cloneUpdateQueue<State>(
 
     firstCapturedEffect: null,
     lastCapturedEffect: null,
+
+    rebaseTime: currentQueue.rebaseTime,
+    rebaseEnd: currentQueue.rebaseEnd,
   };
   return queue;
 }
@@ -446,15 +457,37 @@ export function processUpdateQueue<State>(
   let newBaseState = queue.baseState;
   let newFirstUpdate = null;
   let newExpirationTime = NoWork;
+  let newRebaseTime = NoWork;
+  let newRebaseEnd = null;
+  let prevUpdate = null;
 
   // Iterate through the list of updates to compute the result.
   let update = queue.firstUpdate;
   let resultState = newBaseState;
+
+  // Track whether the update is part of a rebase.
+  // TODO: Should probably split this into two separate loops, instead of
+  // using a boolean.
+  const rebaseTime = queue.rebaseTime;
+  const rebaseEnd = queue.rebaseEnd;
+  let isRebasing = rebaseTime !== NoWork;
+
   while (update !== null) {
     const updateExpirationTime = update.expirationTime;
-    if (updateExpirationTime < renderExpirationTime) {
+    if (prevUpdate === rebaseEnd) {
+      isRebasing = false;
+    }
+    if (
+      // Check if this update should be skipped
+      updateExpirationTime < renderExpirationTime &&
+      // If we're currently rebasing, don't skip this update if we already
+      // committed it.
+      (!rebasedUpdatesNeverUncommit ||
+        (!isRebasing || updateExpirationTime < rebaseTime))
+    ) {
       // This update does not have sufficient priority. Skip it.
-      if (newFirstUpdate === null) {
+      if (newRebaseTime === NoWork) {
+        newRebaseTime = renderExpirationTime;
         // This is the first skipped update. It will be the first update in
         // the new list.
         newFirstUpdate = update;
@@ -501,13 +534,16 @@ export function processUpdateQueue<State>(
       }
     }
     // Continue to the next update.
+    prevUpdate = update;
     update = update.next;
   }
 
   // Separately, iterate though the list of captured updates.
   let newFirstCapturedUpdate = null;
   update = queue.firstCapturedUpdate;
   while (update !== null) {
+    // TODO: Captured updates always have the current render expiration time.
+    // Shouldn't need this priority check, because they will never be skipped.
     const updateExpirationTime = update.expirationTime;
     if (updateExpirationTime < renderExpirationTime) {
       // This update does not have sufficient priority. Skip it.
@@ -565,11 +601,15 @@ export function processUpdateQueue<State>(
     // We processed every update, without skipping. That means the new base
     // state is the same as the result state.
     newBaseState = resultState;
+  } else {
+    newRebaseEnd = prevUpdate;
   }
 
   queue.baseState = newBaseState;
   queue.firstUpdate = newFirstUpdate;
   queue.firstCapturedUpdate = newFirstCapturedUpdate;
+  queue.rebaseEnd = newRebaseEnd;
+  queue.rebaseTime = newRebaseTime;
 
   // Set the remaining expiration time to be whatever is remaining in the queue.
   // This should be fine because the only two other things that contribute to

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

@@ -653,4 +653,134 @@ describe('ReactIncrementalUpdates', () => {
       expect(Scheduler).toFlushAndYield(['Commit: goodbye']);
     });
   });
+
+  it.experimental(
+    'when rebasing, does not exclude updates that were already committed, regardless of priority',
+    async () => {
+      const {useState, useLayoutEffect} = React;
+
+      let pushToLog;
+      function App() {
+        const [log, setLog] = useState('');
+        pushToLog = msg => {
+          setLog(prevLog => prevLog + msg);
+        };
+
+        useLayoutEffect(
+          () => {
+            Scheduler.unstable_yieldValue('Committed: ' + log);
+            if (log === 'B') {
+              // Right after B commits, schedule additional updates.
+              Scheduler.unstable_runWithPriority(
+                Scheduler.unstable_UserBlockingPriority,
+                () => {
+                  pushToLog('C');
+                },
+              );
+              setLog(prevLog => prevLog + 'D');
+            }
+          },
+          [log],
+        );
+
+        return log;
+      }
+
+      const root = ReactNoop.createRoot();
+      await ReactNoop.act(async () => {
+        root.render(<App />);
+      });
+      expect(Scheduler).toHaveYielded(['Committed: ']);
+      expect(root).toMatchRenderedOutput('');
+
+      await ReactNoop.act(async () => {
+        pushToLog('A');
+        Scheduler.unstable_runWithPriority(
+          Scheduler.unstable_UserBlockingPriority,
+          () => {
+            pushToLog('B');
+          },
+        );
+      });
+      expect(Scheduler).toHaveYielded([
+        // A and B are pending. B is higher priority, so we'll render that first.
+        'Committed: B',
+        // Because A comes first in the queue, we're now in rebase mode. B must
+        // be rebased on top of A. Also, in a layout effect, we received two new
+        // updates: C and D. C is user-blocking and D is synchronous.
+        //
+        // First render the synchronous update. What we're testing here is that
+        // B *is not dropped* even though it has lower than sync priority. That's
+        // because we already committed it. However, this render should not
+        // include C, because that update wasn't already committed.
+        'Committed: BD',
+        'Committed: BCD',
+        'Committed: ABCD',
+      ]);
+      expect(root).toMatchRenderedOutput('ABCD');
+    },
+  );
+
+  it.experimental(
+    'when rebasing, does not exclude updates that were already committed, regardless of priority (classes)',
+    async () => {
+      let pushToLog;
+      class App extends React.Component {
+        state = {log: ''};
+        pushToLog = msg => {
+          this.setState(prevState => ({log: prevState.log + msg}));
+        };
+        componentDidUpdate() {
+          Scheduler.unstable_yieldValue('Committed: ' + this.state.log);
+          if (this.state.log === 'B') {
+            // Right after B commits, schedule additional updates.
+            Scheduler.unstable_runWithPriority(
+              Scheduler.unstable_UserBlockingPriority,
+              () => {
+                this.pushToLog('C');
+              },
+            );
+            this.pushToLog('D');
+          }
+        }
+        render() {
+          pushToLog = this.pushToLog;
+          return this.state.log;
+        }
+      }
+
+      const root = ReactNoop.createRoot();
+      await ReactNoop.act(async () => {
+        root.render(<App />);
+      });
+      expect(Scheduler).toHaveYielded([]);
+      expect(root).toMatchRenderedOutput('');
+
+      await ReactNoop.act(async () => {
+        pushToLog('A');
+        Scheduler.unstable_runWithPriority(
+          Scheduler.unstable_UserBlockingPriority,
+          () => {
+            pushToLog('B');
+          },
+        );
+      });
+      expect(Scheduler).toHaveYielded([
+        // A and B are pending. B is higher priority, so we'll render that first.
+        'Committed: B',
+        // Because A comes first in the queue, we're now in rebase mode. B must
+        // be rebased on top of A. Also, in a layout effect, we received two new
+        // updates: C and D. C is user-blocking and D is synchronous.
+        //
+        // First render the synchronous update. What we're testing here is that
+        // B *is not dropped* even though it has lower than sync priority. That's
+        // because we already committed it. However, this render should not
+        // include C, because that update wasn't already committed.
+        'Committed: BD',
+        'Committed: BCD',
+        'Committed: ABCD',
+      ]);
+      expect(root).toMatchRenderedOutput('ABCD');
+    },
+  );
 });

packages/shared/ReactFeatureFlags.js

@@ -92,3 +92,5 @@ export const enableTrustedTypesIntegration = false;
 
 // Flag to turn event.target and event.currentTarget in ReactNative from a reactTag to a component instance
 export const enableNativeTargetAsInstance = false;
+
+export const rebasedUpdatesNeverUncommit = __EXPERIMENTAL__;

packages/shared/forks/ReactFeatureFlags.native-fb.js

@@ -43,6 +43,7 @@ export const warnAboutStringRefs = false;
 export const disableLegacyContext = false;
 export const disableSchedulerTimeoutBasedOnReactExpirationTime = false;
 export const enableTrustedTypesIntegration = false;
+export const rebasedUpdatesNeverUncommit = __EXPERIMENTAL__;
 
 // Only used in www builds.
 export function addUserTimingListener() {

packages/shared/forks/ReactFeatureFlags.native-oss.js

@@ -38,6 +38,7 @@ export const disableLegacyContext = false;
 export const disableSchedulerTimeoutBasedOnReactExpirationTime = false;
 export const enableTrustedTypesIntegration = false;
 export const enableNativeTargetAsInstance = false;
+export const rebasedUpdatesNeverUncommit = __EXPERIMENTAL__;
 
 // Only used in www builds.
 export function addUserTimingListener() {