Render phase updates can now be extracted from the pending queue

packages/react-reconciler/src/ReactFiberHooks.js

@@ -193,11 +193,8 @@ let workInProgressHook: Hook | null = null;
 
 // Whether an update was scheduled during the currently executing render pass.
 let didScheduleRenderPhaseUpdate: boolean = false;
-// Lazily created map of render-phase updates
-let renderPhaseUpdates: Map<
-  UpdateQueue<any, any>,
-  Update<any, any>,
-> | null = null;
+// TODO: Move to custom dispatcher.
+let isRerender = false;
 
 const RE_RENDER_LIMIT = 25;
 
@@ -396,7 +393,6 @@ export function renderWithHooks(
   // workInProgressHook = null;
 
   // didScheduleRenderPhaseUpdate = false;
-  // renderPhaseUpdates = null;
 
   // TODO Warn if no hooks are used at all during mount, then some are used during update.
   // Currently we will identify the update render as a mount because memoizedState === null.
@@ -428,6 +424,8 @@ export function renderWithHooks(
   let children = Component(props, refOrContext);
 
   if (didScheduleRenderPhaseUpdate) {
+    isRerender = true;
+
     // Counter to prevent infinite loops.
     let numberOfReRenders: number = 0;
     do {
@@ -464,7 +462,7 @@ export function renderWithHooks(
       children = Component(props, refOrContext);
     } while (didScheduleRenderPhaseUpdate);
 
-    renderPhaseUpdates = null;
+    isRerender = false;
   }
 
   // We can assume the previous dispatcher is always this one, since we set it
@@ -494,7 +492,6 @@ export function renderWithHooks(
 
   // These were reset above
   // didScheduleRenderPhaseUpdate = false;
-  // renderPhaseUpdates = null;
 
   invariant(
     !didRenderTooFewHooks,
@@ -540,7 +537,7 @@ export function resetHooks(): void {
   }
 
   didScheduleRenderPhaseUpdate = false;
-  renderPhaseUpdates = null;
+  isRerender = false;
 }
 
 function mountWorkInProgressHook(): Hook {
@@ -676,15 +673,17 @@ function updateReducer<S, I, A>(
 
   queue.lastRenderedReducer = reducer;
 
-  if (renderPhaseUpdates !== null) {
+  if (isRerender) {
     // This is a re-render. Apply the new render phase updates to the previous
     // work-in-progress hook.
     const dispatch: Dispatch<A> = (queue.dispatch: any);
-    // Render phase updates are stored in a map of queue -> linked list
-    const firstRenderPhaseUpdate = renderPhaseUpdates.get(queue);
-    if (firstRenderPhaseUpdate !== undefined) {
-      renderPhaseUpdates.delete(queue);
-      let newState = hook.memoizedState;
+    const lastRenderPhaseUpdate = queue.pending;
+    let newState = hook.memoizedState;
+    if (lastRenderPhaseUpdate !== null) {
+      // The queue doesn't persist past this render pass.
+      queue.pending = null;
+
+      const firstRenderPhaseUpdate = lastRenderPhaseUpdate.next;
       let update = firstRenderPhaseUpdate;
       do {
         // Process this render phase update. We don't have to check the
@@ -693,7 +692,7 @@ function updateReducer<S, I, A>(
         const action = update.action;
         newState = reducer(newState, action);
         update = update.next;
-      } while (update !== null);
+      } while (update !== firstRenderPhaseUpdate);
 
       // Mark that the fiber performed work, but only if the new state is
       // different from the current state.
@@ -711,10 +710,8 @@ function updateReducer<S, I, A>(
       }
 
       queue.lastRenderedState = newState;
-
-      return [newState, dispatch];
     }
-    return [hook.memoizedState, dispatch];
+    return [newState, dispatch];
   }
 
   const current: Hook = (currentHook: any);
@@ -1239,80 +1236,56 @@ function dispatchAction<S, A>(
     );
   }
 
-  const alternate = fiber.alternate;
+  const currentTime = requestCurrentTimeForUpdate();
+  const suspenseConfig = requestCurrentSuspenseConfig();
+  const expirationTime = computeExpirationForFiber(
+    currentTime,
+    fiber,
+    suspenseConfig,
+  );
+
+  const update: Update<S, A> = {
+    expirationTime,
+    suspenseConfig,
+    action,
+    eagerReducer: null,
+    eagerState: null,
+    next: (null: any),
+  };
+
+  if (__DEV__) {
+    update.priority = getCurrentPriorityLevel();
+  }
+
+  // Append the update to the end of the list.
+  const pending = queue.pending;
+  if (pending === null) {
+    // This is the first update. Create a circular list.
+    update.next = update;
+  } else {
+    const first = pending.next;
+    if (first !== null) {
+      // Still circular.
+      update.next = first;
+    }
+    pending.next = update;
+  }
+  queue.pending = update;
+
   if (
-    fiber === currentlyRenderingFiber ||
-    (alternate !== null && alternate === currentlyRenderingFiber)
+    currentlyRenderingFiber !== null &&
+    (fiber === currentlyRenderingFiber ||
+      fiber.alternate === currentlyRenderingFiber)
   ) {
     // This is a render phase update. Stash it in a lazily-created map of
     // queue -> linked list of updates. After this render pass, we'll restart
     // and apply the stashed updates on top of the work-in-progress hook.
     didScheduleRenderPhaseUpdate = true;
-    const update: Update<S, A> = {
-      expirationTime: renderExpirationTime,
-      suspenseConfig: null,
-      action,
-      eagerReducer: null,
-      eagerState: null,
-      next: (null: any),
-    };
-    if (__DEV__) {
-      update.priority = getCurrentPriorityLevel();
-    }
-    if (renderPhaseUpdates === null) {
-      renderPhaseUpdates = new Map();
-    }
-    const firstRenderPhaseUpdate = renderPhaseUpdates.get(queue);
-    if (firstRenderPhaseUpdate === undefined) {
-      renderPhaseUpdates.set(queue, update);
-    } else {
-      // Append the update to the end of the list.
-      let lastRenderPhaseUpdate = firstRenderPhaseUpdate;
-      while (lastRenderPhaseUpdate.next !== null) {
-        lastRenderPhaseUpdate = lastRenderPhaseUpdate.next;
-      }
-      lastRenderPhaseUpdate.next = update;
-    }
+    update.expirationTime = renderExpirationTime;
   } else {
-    const currentTime = requestCurrentTimeForUpdate();
-    const suspenseConfig = requestCurrentSuspenseConfig();
-    const expirationTime = computeExpirationForFiber(
-      currentTime,
-      fiber,
-      suspenseConfig,
-    );
-
-    const update: Update<S, A> = {
-      expirationTime,
-      suspenseConfig,
-      action,
-      eagerReducer: null,
-      eagerState: null,
-      next: (null: any),
-    };
-
-    if (__DEV__) {
-      update.priority = getCurrentPriorityLevel();
-    }
-
-    // Append the update to the end of the list.
-    const pending = queue.pending;
-    if (pending === null) {
-      // This is the first update. Create a circular list.
-      update.next = update;
-    } else {
-      const first = pending.next;
-      if (first !== null) {
-        // Still circular.
-        update.next = first;
-      }
-      pending.next = update;
-    }
-    queue.pending = update;
-
     if (
       fiber.expirationTime === NoWork &&
-      (alternate === null || alternate.expirationTime === NoWork)
+      (fiber.alternate === null || fiber.alternate.expirationTime === NoWork)
     ) {
       // The queue is currently empty, which means we can eagerly compute the
       // next state before entering the render phase. If the new state is the