useMemoCache implementation (#25143)

josephsavona · web-flow · commit 3401e9200efb · 2022-09-13T14:44:32.000-07:00
* useMemoCache impl
* test for multiple calls in a component (from custom hook)
* Use array of arrays for multiple calls; use alternate/local as the backup
* code cleanup
* fix internal test
* oops we do not support nullable property access
* Simplify implementation, still have questions on some of the PR feedback though
* Gate all code based on the feature flag
* refactor to use updateQueue
* address feedback
* Try to eliminate size increase in prod bundle
* update to retrigger ci
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';
@@ -177,6 +182,8 @@ type StoreConsistencyCheck<T> = {
 export type FunctionComponentUpdateQueue = {
   lastEffect: Effect | null,
   stores: Array<StoreConsistencyCheck<any>> | null,
+  // NOTE: optional, only set when enableUseMemoCacheHook is enabled
+  memoCache?: MemoCache | null,
 };
 
 type BasicStateAction<S> = (S => S) | S;
@@ -710,10 +717,23 @@ function updateWorkInProgressHook(): Hook {
   return workInProgressHook;
 }
 
-function createFunctionComponentUpdateQueue(): FunctionComponentUpdateQueue {
-  return {
-    lastEffect: null,
-    stores: null,
+// NOTE: defining two versions of this function to avoid size impact when this feature is disabled.
+// Previously this function was inlined, the additional `memoCache` property makes it not inlined.
+let createFunctionComponentUpdateQueue: () => FunctionComponentUpdateQueue;
+if (enableUseMemoCacheHook) {
+  createFunctionComponentUpdateQueue = () => {
+    return {
+      lastEffect: null,
+      stores: null,
+      memoCache: null,
+    };
+  };
+} else {
+  createFunctionComponentUpdateQueue = () => {
+    return {
+      lastEffect: null,
+      stores: null,
+    };
   };
 }
 
@@ -787,7 +807,59 @@ function use<T>(usable: Usable<T>): T {
 }
 
 function useMemoCache(size: number): Array<any> {
-  throw new Error('Not implemented.');
+  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 = currentUpdateQueue.memoCache;
+        if (currentMemoCache != null) {
+          memoCache = {
+            data: currentMemoCache.data.map(array => array.slice()),
+            index: 0,
+          };
+        }
+      }
+    }
+  }
+  // Finally fall back to allocating a fresh instance of the cache
+  if (memoCache == null) {
+    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) {
+    // TODO: consider warning or throwing here
+    if (__DEV__) {
+      console.error(
+        'Expected a constant size argument for each invocation of useMemoCache. ' +
+          'The previous cache was allocated with size %s but size %s was requested.',
+        data.length,
+        size,
+      );
+    }
+  }
+  memoCache.index++;
+  return data;
 }
 
 function basicStateReducer<S>(state: S, action: BasicStateAction<S>): S {
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';
@@ -177,6 +182,8 @@ type StoreConsistencyCheck<T> = {
 export type FunctionComponentUpdateQueue = {
   lastEffect: Effect | null,
   stores: Array<StoreConsistencyCheck<any>> | null,
+  // NOTE: optional, only set when enableUseMemoCacheHook is enabled
+  memoCache?: MemoCache | null,
 };
 
 type BasicStateAction<S> = (S => S) | S;
@@ -710,10 +717,23 @@ function updateWorkInProgressHook(): Hook {
   return workInProgressHook;
 }
 
-function createFunctionComponentUpdateQueue(): FunctionComponentUpdateQueue {
-  return {
-    lastEffect: null,
-    stores: null,
+// NOTE: defining two versions of this function to avoid size impact when this feature is disabled.
+// Previously this function was inlined, the additional `memoCache` property makes it not inlined.
+let createFunctionComponentUpdateQueue: () => FunctionComponentUpdateQueue;
+if (enableUseMemoCacheHook) {
+  createFunctionComponentUpdateQueue = () => {
+    return {
+      lastEffect: null,
+      stores: null,
+      memoCache: null,
+    };
+  };
+} else {
+  createFunctionComponentUpdateQueue = () => {
+    return {
+      lastEffect: null,
+      stores: null,
+    };
   };
 }
 
@@ -787,7 +807,59 @@ function use<T>(usable: Usable<T>): T {
 }
 
 function useMemoCache(size: number): Array<any> {
-  throw new Error('Not implemented.');
+  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 = currentUpdateQueue.memoCache;
+        if (currentMemoCache != null) {
+          memoCache = {
+            data: currentMemoCache.data.map(array => array.slice()),
+            index: 0,
+          };
+        }
+      }
+    }
+  }
+  // Finally fall back to allocating a fresh instance of the cache
+  if (memoCache == null) {
+    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) {
+    // TODO: consider warning or throwing here
+    if (__DEV__) {
+      console.error(
+        'Expected a constant size argument for each invocation of useMemoCache. ' +
+          'The previous cache was allocated with size %s but size %s was requested.',
+        data.length,
+        size,
+      );
+    }
+  }
+  memoCache.index++;
+  return data;
 }
 
 function basicStateReducer<S>(state: S, action: BasicStateAction<S>): S {
diff --git a/packages/react-reconciler/src/ReactInternalTypes.js b/packages/react-reconciler/src/ReactInternalTypes.js
@@ -68,6 +68,11 @@ export type Dependencies = {
   ...
 };
 
+export type MemoCache = {
+  data: Array<Array<any>>,
+  index: number,
+};
+
 // A Fiber is work on a Component that needs to be done or was done. There can
 // be more than one per component.
 export type Fiber = {
diff --git a/packages/react-reconciler/src/__tests__/useMemoCache-test.js b/packages/react-reconciler/src/__tests__/useMemoCache-test.js
