Implement experimental_useOptimisticState

This adds an experimental hook tentatively called useOptimisticState.
(The actual name needs some bikeshedding.)

The headline feature is that you can use it to implement optimistic
updates. If you set some optimistic state during a transition/action,
the state will be automatically reverted once the transition completes.

Another feature is that the optimistic updates will be continually
rebased on top of the latest state.

It's easiest to explain with examples; we'll publish documentation as
the API gets closer to stabilizing. See tests for now.

Technically the use cases for this hook are broader than just optimistic
updates; you could use it implement any sort of "pending" state, such
as the ones exposed by useTransition and useFormStatus. But we expect
people will most often reach for this hook to implement the optimistic
update pattern; simpler cases are covered by those other hooks.

Author: acdlite

packages/react-reconciler/src/ReactFiberHooks.js

@@ -149,11 +149,13 @@ import type {ThenableState} from './ReactFiberThenable';
 import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent';
 import {requestAsyncActionContext} from './ReactFiberAsyncAction';
 import {HostTransitionContext} from './ReactFiberHostContext';
+import {requestTransitionLane} from './ReactFiberRootScheduler';
 
 const {ReactCurrentDispatcher, ReactCurrentBatchConfig} = ReactSharedInternals;
 
 export type Update<S, A> = {
   lane: Lane,
+  revertLane: Lane,
   action: A,
   hasEagerState: boolean,
   eagerState: S | null,
@@ -1136,6 +1138,14 @@ function updateReducer<S, I, A>(
   init?: I => S,
 ): [S, Dispatch<A>] {
   const hook = updateWorkInProgressHook();
+  return updateReducerImpl(hook, ((currentHook: any): Hook), reducer);
+}
+
+function updateReducerImpl<S, A>(
+  hook: Hook,
+  current: Hook,
+  reducer: (S, A) => S,
+): [S, Dispatch<A>] {
   const queue = hook.queue;
 
   if (queue === null) {
@@ -1146,10 +1156,8 @@ function updateReducer<S, I, A>(
 
   queue.lastRenderedReducer = reducer;
 
-  const current: Hook = (currentHook: any);
-
   // The last rebase update that is NOT part of the base state.
-  let baseQueue = current.baseQueue;
+  let baseQueue = hook.baseQueue;
 
   // The last pending update that hasn't been processed yet.
   const pendingQueue = queue.pending;
@@ -1180,7 +1188,7 @@ function updateReducer<S, I, A>(
   if (baseQueue !== null) {
     // We have a queue to process.
     const first = baseQueue.next;
-    let newState = current.baseState;
+    let newState = hook.baseState;
 
     let newBaseState = null;
     let newBaseQueueFirst = null;
@@ -1206,6 +1214,7 @@ function updateReducer<S, I, A>(
         // update/state.
         const clone: Update<S, A> = {
           lane: updateLane,
+          revertLane: update.revertLane,
           action: update.action,
           hasEagerState: update.hasEagerState,
           eagerState: update.eagerState,
@@ -1228,18 +1237,68 @@ function updateReducer<S, I, A>(
       } else {
         // This update does have sufficient priority.
 
-        if (newBaseQueueLast !== null) {
-          const clone: Update<S, A> = {
-            // This update is going to be committed so we never want uncommit
-            // it. Using NoLane works because 0 is a subset of all bitmasks, so
-            // this will never be skipped by the check above.
-            lane: NoLane,
-            action: update.action,
-            hasEagerState: update.hasEagerState,
-            eagerState: update.eagerState,
-            next: (null: any),
-          };
-          newBaseQueueLast = newBaseQueueLast.next = clone;
+        // Check if this is an optimistic update.
+        const revertLane = update.revertLane;
+        if (revertLane === NoLane) {
+          // This is not an optimistic update, and we're going to apply it now.
+          // But, if there were earlier updates that were skipped, we need to
+          // leave this update in the queue so it can be rebased later.
+          if (newBaseQueueLast !== null) {
+            const clone: Update<S, A> = {
+              // This update is going to be committed so we never want uncommit
+              // it. Using NoLane works because 0 is a subset of all bitmasks, so
+              // this will never be skipped by the check above.
+              lane: NoLane,
+              revertLane: NoLane,
+              action: update.action,
+              hasEagerState: update.hasEagerState,
+              eagerState: update.eagerState,
+              next: (null: any),
+            };
+            newBaseQueueLast = newBaseQueueLast.next = clone;
+          }
+        } else {
+          // This is an optimistic update. If the "revert" priority is
+          // sufficient, don't apply the update. Otherwise, apply the update,
+          // but leave it in the queue so it can be either reverted or
+          // rebased in a subsequent render.
+          if (isSubsetOfLanes(renderLanes, revertLane)) {
+            // The transition that this optimistic update is associated with
+            // has finished. Pretend the update doesn't exist by skipping
+            // over it.
+            update = update.next;
+            continue;
+          } else {
+            const clone: Update<S, A> = {
+              // Once we commit an optimistic update, we shouldn't uncommit it
+              // until the transition it is associated with has finished
+              // (represented by revertLane). Using NoLane here works because 0
+              // is a subset of all bitmasks, so this will never be skipped by
+              // the check above.
+              lane: NoLane,
+              // Reuse the same revertLane so we know when the transition
+              // has finished.
+              revertLane: update.revertLane,
+              action: update.action,
+              hasEagerState: update.hasEagerState,
+              eagerState: update.eagerState,
+              next: (null: any),
+            };
+            if (newBaseQueueLast === null) {
+              newBaseQueueFirst = newBaseQueueLast = clone;
+              newBaseState = newState;
+            } else {
+              newBaseQueueLast = newBaseQueueLast.next = clone;
+            }
+            // Update the remaining priority in the queue.
+            // TODO: Don't need to accumulate this. Instead, we can remove
+            // renderLanes from the original lanes.
+            currentlyRenderingFiber.lanes = mergeLanes(
+              currentlyRenderingFiber.lanes,
+              revertLane,
+            );
+            markSkippedUpdateLanes(revertLane);
+          }
         }
 
         // Process this update.
@@ -1899,56 +1958,106 @@ function mountStateImpl<S>(initialState: (() => S) | S): Hook {
     lastRenderedState: (initialState: any),
   };
   hook.queue = queue;
-  const dispatch: Dispatch<BasicStateAction<S>> = (dispatchSetState.bind(
-    null,
-    currentlyRenderingFiber,
-    queue,
-  ): any);
-  queue.dispatch = dispatch;
   return hook;
 }
 
 function mountState<S>(
   initialState: (() => S) | S,
 ): [S, Dispatch<BasicStateAction<S>>] {
   const hook = mountStateImpl(initialState);
-  return [hook.memoizedState, hook.queue.dispatch];
+  const queue = hook.queue;
+  const dispatch: Dispatch<BasicStateAction<S>> = (dispatchSetState.bind(
+    null,
+    currentlyRenderingFiber,
+    queue,
+  ): any);
+  queue.dispatch = dispatch;
+  return [hook.memoizedState, dispatch];
 }
 
 function updateState<S>(
   initialState: (() => S) | S,
 ): [S, Dispatch<BasicStateAction<S>>] {
-  return updateReducer(basicStateReducer, (initialState: any));
+  return updateReducer(basicStateReducer, initialState);
 }
 
 function rerenderState<S>(
   initialState: (() => S) | S,
 ): [S, Dispatch<BasicStateAction<S>>] {
-  return rerenderReducer(basicStateReducer, (initialState: any));
+  return rerenderReducer(basicStateReducer, initialState);
 }
 
 function mountOptimisticState<S, A>(
   passthrough: S,
   reducer: ?(S, A) => S,
 ): [S, (A) => void] {
-  // $FlowFixMe - TODO: Actual implementation
-  return mountState(passthrough);
+  const hook = mountWorkInProgressHook();
+  hook.memoizedState = hook.baseState = passthrough;
+  const queue: UpdateQueue<S, A> = {
+    pending: null,
+    lanes: NoLanes,
+    dispatch: null,
+    // Optimistic state does not use the eager update optimization.
+    lastRenderedReducer: null,
+    lastRenderedState: null,
+  };
+  hook.queue = queue;
+  // This is different than the normal setState function.
+  const dispatch: A => void = (dispatchOptimisticSetState.bind(
+    null,
+    currentlyRenderingFiber,
+    true,
+    queue,
+  ): any);
+  queue.dispatch = dispatch;
+  return [passthrough, dispatch];
 }
 
 function updateOptimisticState<S, A>(
   passthrough: S,
   reducer: ?(S, A) => S,
 ): [S, (A) => void] {
-  // $FlowFixMe - TODO: Actual implementation
-  return updateState(passthrough);
+  const hook = updateWorkInProgressHook();
+
+  // Optimistic updates are always rebased on top of the latest value passed in
+  // as an argument. It's called a passthrough because if there are no pending
+  // updates, it will be returned as-is.
+  //
+  // Reset the base state and memoized state to the passthrough. Future
+  // updates will be applied on top of this.
+  hook.baseState = hook.memoizedState = passthrough;
+
+  // If a reducer is not provided, default to the same one used by useState.
+  const resolvedReducer: (S, A) => S =
+    typeof reducer === 'function' ? reducer : (basicStateReducer: any);
+
+  return updateReducerImpl(hook, ((currentHook: any): Hook), resolvedReducer);
 }
 
 function rerenderOptimisticState<S, A>(
   passthrough: S,
   reducer: ?(S, A) => S,
 ): [S, (A) => void] {
-  // $FlowFixMe - TODO: Actual implementation
-  return rerenderState(passthrough);
+  // Unlike useState, useOptimisticState doesn't support render phase updates.
+  // Also unlike useState, we need to replay all pending updates again in case
+  // the passthrough value changed.
+  //
+  // So instead of a forked re-render implementation that knows how to handle
+  // render phase udpates, we can use the same implementation as during a
+  // regular mount or update.
+
+  if (currentHook !== null) {
+    // This is an update. Process the update queue.
+    return updateOptimisticState(passthrough, reducer);
+  }
+
+  // This is a mount. No updates to process.
+  const hook = updateWorkInProgressHook();
+  // Reset the base state and memoized state to the passthrough. Future
+  // updates will be applied on top of this.
+  hook.baseState = hook.memoizedState = passthrough;
+  const dispatch = hook.queue.dispatch;
+  return [passthrough, dispatch];
 }
 
 function pushEffect(
@@ -2490,9 +2599,15 @@ function startTransition<S>(
     higherEventPriority(previousPriority, ContinuousEventPriority),
   );
 
+  // We don't really need to use an optimistic update here, because we schedule
+  // a second "revert" update below (which we use to suspend the transition
+  // until the async action scope has finished). But we'll use an optimistic
+  // update anyway to make it less likely the behavior accidentally diverges;
+  // for example, both an optimistic update and this one should share the
+  // same lane.
+  dispatchOptimisticSetState(fiber, false, queue, pendingState);
+
   const prevTransition = ReactCurrentBatchConfig.transition;
-  ReactCurrentBatchConfig.transition = null;
-  dispatchSetState(fiber, queue, pendingState);
   const currentTransition = (ReactCurrentBatchConfig.transition =
     ({}: BatchConfigTransition));
 
@@ -2827,6 +2942,7 @@ function dispatchReducerAction<S, A>(
 
   const update: Update<S, A> = {
     lane,
+    revertLane: NoLane,
     action,
     hasEagerState: false,
     eagerState: null,
@@ -2865,6 +2981,7 @@ function dispatchSetState<S, A>(
 
   const update: Update<S, A> = {
     lane,
+    revertLane: NoLane,
     action,
     hasEagerState: false,
     eagerState: null,
@@ -2928,6 +3045,54 @@ function dispatchSetState<S, A>(
   markUpdateInDevTools(fiber, lane, action);
 }
 
+function dispatchOptimisticSetState<S, A>(
+  fiber: Fiber,
+  throwIfDuringRender: boolean,
+  queue: UpdateQueue<S, A>,
+  action: A,
+): void {
+  const update: Update<S, A> = {
+    // An optimistic update commits synchronously.
+    lane: SyncLane,
+    // After committing, the optimistic update is "reverted" using the same
+    // lane as the transition it's associated with.
+    //
+    // TODO: Warn if there's no transition/action associated with this
+    // optimistic update.
+    revertLane: requestTransitionLane(),
+    action,
+    hasEagerState: false,
+    eagerState: null,
+    next: (null: any),
+  };
+
+  if (isRenderPhaseUpdate(fiber)) {
+    // When calling startTransition during render, this warns instead of
+    // throwing because throwing would be a breaking change. setOptimisticState
+    // is a new API so it's OK to throw.
+    if (throwIfDuringRender) {
+      throw new Error('Cannot update optimistic state while rendering.');
+    } else {
+      // startTransition was called during render. We don't need to do anything
+      // besides warn here because the render phase update would be overidden by
+      // the second update, anyway. We can remove this branch and make it throw
+      // in a future release.
+      if (__DEV__) {
+        console.error('Cannot call startTransition state while rendering.');
+      }
+    }
+  } else {
+    const root = enqueueConcurrentHookUpdate(fiber, queue, update, SyncLane);
+    if (root !== null) {
+      scheduleUpdateOnFiber(root, fiber, SyncLane);
+      // Optimistic updates are always synchronous, so we don't need to call
+      // entangleTransitionUpdate here.
+    }
+  }
+
+  markUpdateInDevTools(fiber, SyncLane, action);
+}
+
 function isRenderPhaseUpdate(fiber: Fiber): boolean {
   const alternate = fiber.alternate;
   return (