refactor to use updateQueue

facebook · Sep 8, 2022 · bebca2f · bebca2f
1 parent 22f89ed
commit bebca2f
Show file tree

Hide file tree

Showing 8 changed files with 89 additions and 120 deletions.
diff --git a/packages/react-reconciler/src/ReactFiber.new.js b/packages/react-reconciler/src/ReactFiber.new.js
@@ -32,7 +32,6 @@ import {
   allowConcurrentByDefault,
   enableTransitionTracing,
   enableDebugTracing,
-  enableUseMemoCacheHook,
 } from 'shared/ReactFeatureFlags';
 import {NoFlags, Placement, StaticMask} from './ReactFiberFlags';
 import {ConcurrentRoot} from './ReactRootTags';
@@ -145,9 +144,6 @@ function FiberNode(
   this.updateQueue = null;
   this.memoizedState = null;
   this.dependencies = null;
-  if (enableUseMemoCacheHook) {
-    this.memoCache = null;
-  }
 
   this.mode = mode;
 
@@ -314,9 +310,6 @@ export function createWorkInProgress(current: Fiber, pendingProps: any): Fiber {
   workInProgress.memoizedProps = current.memoizedProps;
   workInProgress.memoizedState = current.memoizedState;
   workInProgress.updateQueue = current.updateQueue;
-  if (enableUseMemoCacheHook) {
-    workInProgress.memoCache = current.memoCache;
-  }
 
   // Clone the dependencies object. This is mutated during the render phase, so
   // it cannot be shared with the current fiber.
@@ -858,9 +851,6 @@ export function assignFiberPropertiesInDEV(
   target.pendingProps = source.pendingProps;
   target.memoizedProps = source.memoizedProps;
   target.updateQueue = source.updateQueue;
-  if (enableUseMemoCacheHook) {
-    target.memoCache = source.memoCache;
-  }
   target.memoizedState = source.memoizedState;
   target.dependencies = source.dependencies;
   target.mode = source.mode;

diff --git a/packages/react-reconciler/src/ReactFiber.old.js b/packages/react-reconciler/src/ReactFiber.old.js
@@ -32,7 +32,6 @@ import {
   allowConcurrentByDefault,
   enableTransitionTracing,
   enableDebugTracing,
-  enableUseMemoCacheHook,
 } from 'shared/ReactFeatureFlags';
 import {NoFlags, Placement, StaticMask} from './ReactFiberFlags';
 import {ConcurrentRoot} from './ReactRootTags';
@@ -145,9 +144,6 @@ function FiberNode(
   this.updateQueue = null;
   this.memoizedState = null;
   this.dependencies = null;
-  if (enableUseMemoCacheHook) {
-    this.memoCache = null;
-  }
 
   this.mode = mode;
 
@@ -314,9 +310,6 @@ export function createWorkInProgress(current: Fiber, pendingProps: any): Fiber {
   workInProgress.memoizedProps = current.memoizedProps;
   workInProgress.memoizedState = current.memoizedState;
   workInProgress.updateQueue = current.updateQueue;
-  if (enableUseMemoCacheHook) {
-    workInProgress.memoCache = current.memoCache;
-  }
 
   // Clone the dependencies object. This is mutated during the render phase, so
   // it cannot be shared with the current fiber.
@@ -858,9 +851,6 @@ export function assignFiberPropertiesInDEV(
   target.pendingProps = source.pendingProps;
   target.memoizedProps = source.memoizedProps;
   target.updateQueue = source.updateQueue;
-  if (enableUseMemoCacheHook) {
-    target.memoCache = source.memoCache;
-  }
   target.memoizedState = source.memoizedState;
   target.dependencies = source.dependencies;
   target.mode = source.mode;

diff --git a/packages/react-reconciler/src/ReactFiberHooks.new.js b/packages/react-reconciler/src/ReactFiberHooks.new.js
@@ -16,7 +16,12 @@ import type {
   Usable,
   Thenable,
 } from 'shared/ReactTypes';
-import type {Fiber, Dispatcher, HookType} from './ReactInternalTypes';
+import type {
+  Fiber,
+  Dispatcher,
+  HookType,
+  MemoCache,
+} from './ReactInternalTypes';
 import type {Lanes, Lane} from './ReactFiberLane.new';
 import type {HookFlags} from './ReactHookEffectTags';
 import type {FiberRoot} from './ReactInternalTypes';
@@ -173,6 +178,7 @@ type StoreConsistencyCheck<T> = {|
 export type FunctionComponentUpdateQueue = {|
   lastEffect: Effect | null,
   stores: Array<StoreConsistencyCheck<any>> | null,
+  memoCache: MemoCache | null,
 |};
 
 type BasicStateAction<S> = (S => S) | S;
@@ -395,20 +401,6 @@ export function renderWithHooks<Props, SecondArg>(
       current !== null && current.type !== workInProgress.type;
   }
 
-  // Reset the memoCache index and create a backup copy in case of a setState during render
-  // or error, either of which can leave the cache in an inconsistent state
-  let previousMemoCache = null;
-  if (enableUseMemoCacheHook) {
-    previousMemoCache = workInProgress.memoCache;
-    if (previousMemoCache !== null) {
-      const memoCache = {
-        data: previousMemoCache.data.map(array => array.slice()),
-        index: 0,
-      };
-      workInProgress.memoCache = memoCache;
-    }
-  }
-
   workInProgress.memoizedState = null;
   workInProgress.updateQueue = null;
   workInProgress.lanes = NoLanes;
@@ -480,17 +472,6 @@ export function renderWithHooks<Props, SecondArg>(
 
       workInProgress.updateQueue = null;
 
-      if (enableUseMemoCacheHook) {
-        if (previousMemoCache !== null) {
-          // Setting state during render could leave the cache in an inconsistent state,
-          // reset to the previous state before re-rendering.
-          workInProgress.memoCache = {
-            data: previousMemoCache.data.map(array => array.slice()),
-            index: 0,
-          };
-        }
-      }
-
       if (__DEV__) {
         // Also validate hook order for cascading updates.
         hookTypesUpdateIndexDev = -1;
@@ -604,7 +585,7 @@ export function bailoutHooks(
   current.lanes = removeLanes(current.lanes, lanes);
 }
 
-export function resetHooksAfterThrow(erroredWork: Fiber | null): void {
+export function resetHooksAfterThrow(): void {
   // We can assume the previous dispatcher is always this one, since we set it
   // at the beginning of the render phase and there's no re-entrance.
   ReactCurrentDispatcher.current = ContextOnlyDispatcher;
@@ -629,18 +610,6 @@ export function resetHooksAfterThrow(erroredWork: Fiber | null): void {
     didScheduleRenderPhaseUpdate = false;
   }
 
-  if (enableUseMemoCacheHook) {
-    if (erroredWork != null) {
-      const memoCache = erroredWork.memoCache;
-      if (memoCache !== null) {
-        // The memo cache may be in an inconsistent state, reset to the version from
-        // the alternate if available, or clear the cache completely.
-        const alternate = erroredWork.alternate;
-        erroredWork.memoCache = alternate != null ? alternate.memoCache : null;
-      }
-    }
-  }
-
   renderLanes = NoLanes;
   currentlyRenderingFiber = (null: any);
 
@@ -747,6 +716,7 @@ function createFunctionComponentUpdateQueue(): FunctionComponentUpdateQueue {
   return {
     lastEffect: null,
     stores: null,
+    memoCache: null,
   };
 }
 
@@ -818,22 +788,48 @@ function use<T>(usable: Usable<T>): T {
 }
 
 function useMemoCache(size: number): Array<any> {
-  let memoCache = currentlyRenderingFiber.memoCache;
+  let memoCache = null;
+  // Fast-path, load memo cache from wip fiber if already prepared
+  let updateQueue: FunctionComponentUpdateQueue | null = (currentlyRenderingFiber.updateQueue: any);
+  if (updateQueue !== null) {
+    memoCache = updateQueue.memoCache;
+  }
+  // Otherwise clone from the current fiber
+  // TODO: not sure how to access the current fiber here other than going through
+  // currentlyRenderingFiber.alternate
+  if (memoCache === null) {
+    const current: Fiber | null = currentlyRenderingFiber.alternate;
+    if (current !== null) {
+      const currentUpdateQueue: FunctionComponentUpdateQueue | null = (current.updateQueue: any);
+      if (currentUpdateQueue !== null) {
+        const currentMemoCache: MemoCache | null = currentUpdateQueue.memoCache;
+        if (currentMemoCache !== null) {
+          memoCache = {
+            data: currentMemoCache.data.map(array => array.slice()),
+            index: 0,
+          };
+        }
+      }
+    }
+  }
+  // Finally fall back to allocating a fresh instance
   if (memoCache === null) {
-    memoCache = currentlyRenderingFiber.memoCache = {
+    memoCache = {
       data: [],
       index: 0,
     };
   }
+  if (updateQueue === null) {
+    updateQueue = createFunctionComponentUpdateQueue();
+    currentlyRenderingFiber.updateQueue = updateQueue;
+  }
+  updateQueue.memoCache = memoCache;
+
   let data = memoCache.data[memoCache.index];
   if (data === undefined) {
     data = memoCache.data[memoCache.index] = new Array(size);
   } else if (data.length !== size) {
-    if (__DEV__) {
-      console.warn(
-        'Expected each useMemoCache to receive a consistent size argument, found different sizes',
-      );
-    }
+    // TODO: consider warning or throwing here
   }
   memoCache.index++;
   return data;

diff --git a/packages/react-reconciler/src/ReactFiberHooks.old.js b/packages/react-reconciler/src/ReactFiberHooks.old.js
@@ -16,7 +16,12 @@ import type {
   Usable,
   Thenable,
 } from 'shared/ReactTypes';
-import type {Fiber, Dispatcher, HookType} from './ReactInternalTypes';
+import type {
+  Fiber,
+  Dispatcher,
+  HookType,
+  MemoCache,
+} from './ReactInternalTypes';
 import type {Lanes, Lane} from './ReactFiberLane.old';
 import type {HookFlags} from './ReactHookEffectTags';
 import type {FiberRoot} from './ReactInternalTypes';
@@ -173,6 +178,7 @@ type StoreConsistencyCheck<T> = {|
 export type FunctionComponentUpdateQueue = {|
   lastEffect: Effect | null,
   stores: Array<StoreConsistencyCheck<any>> | null,
+  memoCache: MemoCache | null,
 |};
 
 type BasicStateAction<S> = (S => S) | S;
@@ -395,20 +401,6 @@ export function renderWithHooks<Props, SecondArg>(
       current !== null && current.type !== workInProgress.type;
   }
 
-  // Reset the memoCache index and create a backup copy in case of a setState during render
-  // or error, either of which can leave the cache in an inconsistent state
-  let previousMemoCache = null;
-  if (enableUseMemoCacheHook) {
-    previousMemoCache = workInProgress.memoCache;
-    if (previousMemoCache !== null) {
-      const memoCache = {
-        data: previousMemoCache.data.map(array => array.slice()),
-        index: 0,
-      };
-      workInProgress.memoCache = memoCache;
-    }
-  }
-
   workInProgress.memoizedState = null;
   workInProgress.updateQueue = null;
   workInProgress.lanes = NoLanes;
@@ -480,17 +472,6 @@ export function renderWithHooks<Props, SecondArg>(
 
       workInProgress.updateQueue = null;
 
-      if (enableUseMemoCacheHook) {
-        if (previousMemoCache !== null) {
-          // Setting state during render could leave the cache in an inconsistent state,
-          // reset to the previous state before re-rendering.
-          workInProgress.memoCache = {
-            data: previousMemoCache.data.map(array => array.slice()),
-            index: 0,
-          };
-        }
-      }
-
       if (__DEV__) {
         // Also validate hook order for cascading updates.
         hookTypesUpdateIndexDev = -1;
@@ -604,7 +585,7 @@ export function bailoutHooks(
   current.lanes = removeLanes(current.lanes, lanes);
 }
 
-export function resetHooksAfterThrow(erroredWork: Fiber | null): void {
+export function resetHooksAfterThrow(): void {
   // We can assume the previous dispatcher is always this one, since we set it
   // at the beginning of the render phase and there's no re-entrance.
   ReactCurrentDispatcher.current = ContextOnlyDispatcher;
@@ -629,18 +610,6 @@ export function resetHooksAfterThrow(erroredWork: Fiber | null): void {
     didScheduleRenderPhaseUpdate = false;
   }
 
-  if (enableUseMemoCacheHook) {
-    if (erroredWork != null) {
-      const memoCache = erroredWork.memoCache;
-      if (memoCache !== null) {
-        // The memo cache may be in an inconsistent state, reset to the version from
-        // the alternate if available, or clear the cache completely.
-        const alternate = erroredWork.alternate;
-        erroredWork.memoCache = alternate != null ? alternate.memoCache : null;
-      }
-    }
-  }
-
   renderLanes = NoLanes;
   currentlyRenderingFiber = (null: any);
 
@@ -747,6 +716,7 @@ function createFunctionComponentUpdateQueue(): FunctionComponentUpdateQueue {
   return {
     lastEffect: null,
     stores: null,
+    memoCache: null,
   };
 }
 
@@ -818,22 +788,48 @@ function use<T>(usable: Usable<T>): T {
 }
 
 function useMemoCache(size: number): Array<any> {
-  let memoCache = currentlyRenderingFiber.memoCache;
+  let memoCache = null;
+  // Fast-path, load memo cache from wip fiber if already prepared
+  let updateQueue: FunctionComponentUpdateQueue | null = (currentlyRenderingFiber.updateQueue: any);
+  if (updateQueue !== null) {
+    memoCache = updateQueue.memoCache;
+  }
+  // Otherwise clone from the current fiber
+  // TODO: not sure how to access the current fiber here other than going through
+  // currentlyRenderingFiber.alternate
+  if (memoCache === null) {
+    const current: Fiber | null = currentlyRenderingFiber.alternate;
+    if (current !== null) {
+      const currentUpdateQueue: FunctionComponentUpdateQueue | null = (current.updateQueue: any);
+      if (currentUpdateQueue !== null) {
+        const currentMemoCache: MemoCache | null = currentUpdateQueue.memoCache;
+        if (currentMemoCache !== null) {
+          memoCache = {
+            data: currentMemoCache.data.map(array => array.slice()),
+            index: 0,
+          };
+        }
+      }
+    }
+  }
+  // Finally fall back to allocating a fresh instance
   if (memoCache === null) {
-    memoCache = currentlyRenderingFiber.memoCache = {
+    memoCache = {
       data: [],
       index: 0,
     };
   }
+  if (updateQueue === null) {
+    updateQueue = createFunctionComponentUpdateQueue();
+    currentlyRenderingFiber.updateQueue = updateQueue;
+  }
+  updateQueue.memoCache = memoCache;
+
   let data = memoCache.data[memoCache.index];
   if (data === undefined) {
     data = memoCache.data[memoCache.index] = new Array(size);
   } else if (data.length !== size) {
-    if (__DEV__) {
-      console.warn(
-        'Expected each useMemoCache to receive a consistent size argument, found different sizes',
-      );
-    }
+    // TODO: consider warning or throwing here
   }
   memoCache.index++;
   return data;

diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.new.js b/packages/react-reconciler/src/ReactFiberWorkLoop.new.js
@@ -1691,7 +1691,7 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
 function handleThrow(root, thrownValue): void {
   // Reset module-level state that was set during the render phase.
   resetContextDependencies();
-  resetHooksAfterThrow(workInProgress);
+  resetHooksAfterThrow();
   resetCurrentDebugFiberInDEV();
   // TODO: I found and added this missing line while investigating a
   // separate issue. Write a regression test using string refs.
@@ -3307,7 +3307,7 @@ if (__DEV__ && replayFailedUnitOfWorkWithInvokeGuardedCallback) {
       // Keep this code in sync with handleThrow; any changes here must have
       // corresponding changes there.
       resetContextDependencies();
-      resetHooksAfterThrow(unitOfWork);
+      resetHooksAfterThrow();
       // Don't reset current debug fiber, since we're about to work on the
       // same fiber again.