Mark range of updates that need to be rebased

To enforce that updates that are committed can never be un-committed,
even during a rebase, we need to track which updates are part of the
rebase. The current implementation doesn't do this properly. It has a
hidden assumption that, when rebasing, the next `renderExpirationTime`
will represent a later expiration time that the original one. That's
usually true, but it's not *always* true: there could be another update
at even higher priority.

This requires two extra fields on the update queue. I really don't like
that the update queue logic has gotten even more complicated. It's
tempting to say that we should remove rebasing entirely, and that update
queues must always be updated at the same priority. But I'm hesitant to
jump to that conclusion — rebasing can be confusing in the edge cases,
but it's also convenient. Enforcing single-priority queues would really
hurt the ergonomics of useReducer, for example, where multiple values
are centralized in a single queue. It especially hurts the ergonomics
of classes, where there's only a single queue per class.

So it's something to think about, but I don't think "no more rebasing"
is an acceptable answer on its own.

Author: acdlite

packages/react-reconciler/src/ReactFiberHooks.js

@@ -144,8 +144,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 +584,8 @@ function mountWorkInProgressHook(): Hook {
     baseState: null,
     queue: null,
     baseUpdate: null,
+    rebaseEnd: null,
+    rebaseTime: NoWork,
 
     next: null,
   };
@@ -621,6 +627,8 @@ function updateWorkInProgressHook(): Hook {
       baseState: currentHook.baseState,
       queue: currentHook.queue,
       baseUpdate: currentHook.baseUpdate,
+      rebaseEnd: currentHook.rebaseEnd,
+      rebaseTime: currentHook.rebaseTime,
 
       next: null,
     };
@@ -735,8 +743,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 +765,35 @@ 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 {
+      if (prevUpdate === rebaseEnd) {
+        isRebasing = false;
+      }
       const updateExpirationTime = update.expirationTime;
-      if (updateExpirationTime < renderExpirationTime) {
+      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.
+        (!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 +813,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 +827,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 +841,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

@@ -138,6 +138,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 +175,8 @@ export function createUpdateQueue<State>(baseState: State): UpdateQueue<State> {
     lastEffect: null,
     firstCapturedEffect: null,
     lastCapturedEffect: null,
+    rebaseTime: NoWork,
+    rebaseEnd: null,
   };
   return queue;
 }
@@ -194,6 +199,9 @@ function cloneUpdateQueue<State>(
 
     firstCapturedEffect: null,
     lastCapturedEffect: null,
+
+    rebaseTime: currentQueue.rebaseTime,
+    rebaseEnd: currentQueue.rebaseEnd,
   };
   return queue;
 }
@@ -446,15 +454,36 @@ 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.
+      (!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 +530,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 +597,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