Repository: facebook/react
Pull request: useMemoCache implementation

## Summary

context: I'm still a noob to this codebase so this started as a hacky implementation to get feedback. 

Initial implementation of useMemoCache that aims to at least get the basics correct, even if it doesn't handle all cases and/or there are better ways to implement it (feedback please!). The approach adds a new `memoCache: MemoCache | null` property on Fiber. MemoCache contains a current version of the cache and an optional previous version: the previous version is used to restore in the case of a setState during render and/or an error during render. In both cases the cache could become invalid (only partially updated), so restoring in these cases ensures the cache is always in a consistent state. To avoid corrupting the "previous" cache, it is cloned prior to rendering.

I wasn't sure if it would be safe to use the alternate fiber's memoCache as the previous instance — both the main and alternate fiber could be rendered concurrently (right?) in which case they could fail and need to restore independently. Let me know if this assumption is incorrect.

## How did you test this change?

New unit test. 

Changed files:
filename: packages/react-reconciler/src/ReactFiberHooks.new.js
patch: @@ -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 {

filename: packages/react-reconciler/src/ReactFiberHooks.old.js
patch: @@ -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 {

filename: packages/react-reconciler/src/ReactInternalTypes.js
patch: @@ -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 = {

filename: packages/react-reconciler/src/__tests__/useMemoCache-test.js
patch: @@ -0,0 +1,367 @@
+let React;
+let ReactNoop;
+let act;
+let useState;
+let useMemoCache;
+let ErrorBoundary;
+
+describe('useMemoCache()', () => {
+  beforeEach(() => {
+    jest.resetModules();
+
+    React = require('react');
+    ReactNoop = require('react-noop-renderer');
+    act = require('jest-react').act;
+    useState = React.useState;
+    useMemoCache = React.unstable_useMemoCache;
+
+    class _ErrorBoundary extends React.Component {
+      constructor(props) {
+        super(props);
+        this.state = {hasError: false};
+      }
+
+      static getDerivedStateFromError(error) {
+        // Update state so the next render will show the fallback UI.
+        return {hasError: true};
+      }
+
+      componentDidCatch(error, errorInfo) {}
+
+      render() {
+        if (this.state.hasError) {
+          // You can render any custom fallback UI
+          return <h1>Something went wrong.</h1>;
+        }
+
+        return this.props.children;
+      }
+    }
+    ErrorBoundary = _ErrorBoundary;
+  });
+
+  // @gate enableUseMemoCacheHook
+  test('render component using cache', async () => {
+    function Component(props) {
+      const cache = useMemoCache(1);
+      expect(Array.isArray(cache)).toBe(true);
+      expect(cache.length).toBe(1);
+      expect(cache[0]).toBe(undefined);
+      return 'Ok';
+    }
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      root.render(<Component />);
+    });
+    expect(root).toMatchRenderedOutput('Ok');
+  });
+
+  // @gate enableUseMemoCacheHook
+  test('update component using cache', async () => {
+    let setX;
+    let forceUpdate;
+    function Component(props) {
+      const cache = useMemoCache(4);
+
+      // x is used to produce a `data` object passed to the child
+      const [x, _setX] = useState(0);
+      setX = _setX;
+      const c_x = x !== cache[0];
+      cache[0] = x;
+
+      // n is passed as-is to the child as a cache breaker
+      const [n, setN] = useState(0);
+      forceUpdate = () => setN(a => a + 1);
+      const c_n = n !== cache[1];
+      cache[1] = n;
+
+      let data;
+      if (c_x) {
+        data = cache[2] = {text: `Count ${x}`};
+      } else {
+        data = cache[2];
+      }
+      if (c_x || c_n) {
+        return (cache[3] = <Text data={data} n={n} />);
+      } else {
+        return cache[3];
+      }
+    }
+    let data;
+    const Text = jest.fn(function Text(props) {
+      data = props.data;
+      return data.text;
+    });
+
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      root.render(<Component />);
+    });
+    expect(root).toMatchRenderedOutput('Count 0');
+    expect(Text).toBeCalledTimes(1);
+    const data0 = data;
+
+    // Changing x should reset the data object
+    await act(async () => {
+      setX(1);
+    });
+    expect(root).toMatchRenderedOutput('Count 1');
+    expect(Text).toBeCalledTimes(2);
+    expect(data).not.toBe(data0);
+    const data1 = data;
+
+    // Forcing an unrelated update shouldn't recreate the
+    // data object.
+    await act(async () => {
+      forceUpdate();
+    });
+    expect(root).toMatchRenderedOutput('Count 1');
+    expect(Text).toBeCalledTimes(3);
+    expect(data).toBe(data1); // confirm that the cache persisted across renders
+  });
+
+  // @gate enableUseMemoCacheHook
+  test('update component using cache with setstate during render', async () => {
+    let setX;
+    let setN;
+    function Component(props) {
+      const cache = useMemoCache(4);
+
+      // x is used to produce a `data` object passed to the child
+      const [x, _setX] = useState(0);
+      setX = _setX;
+      const c_x = x !== cache[0];
+      cache[0] = x;
+
+      // n is passed as-is to the child as a cache breaker
+      const [n, _setN] = useState(0);
+      setN = _setN;
+      const c_n = n !== cache[1];
+      cache[1] = n;
+
+      // NOTE: setstate and early return here means that x will update
+      // without the data value being updated. Subsequent renders could
+      // therefore think that c_x = false (hasn't changed) and skip updating
+      // data.
+      // The memoizing compiler will have to handle this case, but the runtime
+      // can help by falling back to resetting the cache if a setstate occurs
+      // during render (this mirrors what we do for useMemo and friends)
+      if (n === 1) {
+        setN(2);
+        return;
+      }
+
+      let data;
+      if (c_x) {
+        data = cache[2] = {text: `Count ${x}`};
+      } else {
+        data = cache[2];
+      }
+      if (c_x || c_n) {
+        return (cache[3] = <Text data={data} n={n} />);
+      } else {
+        return cache[3];
+      }
+    }
+    let data;
+    const Text = jest.fn(function Text(props) {
+      data = props.data;
+      return data.text;
+    });
+
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      root.render(<Component />);
+    });
+    expect(root).toMatchRenderedOutput('Count 0');
+    expect(Text).toBeCalledTimes(1);
+    const data0 = data;
+
+    // Simultaneously trigger an update to x (should create a new data value)
+    // and trigger the setState+early return. The runtime should reset the cache
+    // to avoid an inconsistency
+    await act(async () => {
+      setX(1);
+      setN(1);
+    });
+    expect(root).toMatchRenderedOutput('Count 1');
+    expect(Text).toBeCalledTimes(2);
+    expect(data).not.toBe(data0);
+    const data1 = data;
+
+    // Forcing an unrelated update shouldn't recreate the
+    // data object.
+    await act(async () => {
+      setN(3);
+    });
+    expect(root).toMatchRenderedOutput('Count 1');
+    expect(Text).toBeCalledTimes(3);
+    expect(data).toBe(data1); // confirm that the cache persisted across renders
+  });
+
+  // @gate enableUseMemoCacheHook
+  test('update component using cache with throw during render', async () => {
+    let setX;
+    let setN;
+    let shouldFail = true;
+    function Component(props) {
+      const cache = useMemoCache(4);
+
+      // x is used to produce a `data` object passed to the child
+      const [x, _setX] = useState(0);
+      setX = _setX;
+      const c_x = x !== cache[0];
+      cache[0] = x;
+
+      // n is passed as-is to the child as a cache breaker
+      const [n, _setN] = useState(0);
+      setN = _setN;
+      const c_n = n !== cache[1];
+      cache[1] = n;
+
+      // NOTE the initial failure will trigger a re-render, after which the function
+      // will early return. This validates that the runtime resets the cache on error:
+      // if it doesn't the cache will be corrupt, with the cached version of data
+      // out of data from the cached version of x.
+      if (n === 1) {
+        if (shouldFail) {
+          shouldFail = false;
+          throw new Error('failed');
+        }
+        setN(2);
+        return;
+      }
+
+      let data;
+      if (c_x) {
+        data = cache[2] = {text: `Count ${x}`};
+      } else {
+        data = cache[2];
+      }
+      if (c_x || c_n) {
+        return (cache[3] = <Text data={data} n={n} />);
+      } else {
+        return cache[3];
+      }
+    }
+    let data;
+    const Text = jest.fn(function Text(props) {
+      data = props.data;
+      return data.text;
+    });
+
+    spyOnDev(console, 'error');
+
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      root.render(
+        <ErrorBoundary>
+          <Component />
+        </ErrorBoundary>,
+      );
+    });
+    expect(root).toMatchRenderedOutput('Count 0');
+    expect(Text).toBeCalledTimes(1);
+    const data0 = data;
+
+    // Simultaneously trigger an update to x (should create a new data value)
+    // and trigger the setState+early return. The runtime should reset the cache
+    // to avoid an inconsistency
+    await act(async () => {
+      // this update bumps the count
+      setX(1);
+      // this triggers a throw.
+      setN(1);
+    });
+    expect(root).toMatchRenderedOutput('Count 1');
+    expect(Text).toBeCalledTimes(2);
+    expect(data).not.toBe(data0);
+    const data1 = data;
+
+    // Forcing an unrelated update shouldn't recreate the
+    // data object.
+    await act(async () => {
+      setN(3);
+    });
+    expect(root).toMatchRenderedOutput('Count 1');
+    expect(Text).toBeCalledTimes(3);
+    expect(data).toBe(data1); // confirm that the cache persisted across renders
+  });
+
+  // @gate enableUseMemoCacheHook
+  test('update component and custom hook with caches', async () => {
+    let setX;
+    let forceUpdate;
+    function Component(props) {
+      const cache = useMemoCache(4);
+
+      // x is used to produce a `data` object passed to the child
+      const [x, _setX] = useState(0);
+      setX = _setX;
+      const c_x = x !== cache[0];
+      cache[0] = x;
+
+      // n is passed as-is to the child as a cache breaker
+      const [n, setN] = useState(0);
+      forceUpdate = () => setN(a => a + 1);
+      const c_n = n !== cache[1];
+      cache[1] = n;
+
+      let _data;
+      if (c_x) {
+        _data = cache[2] = {text: `Count ${x}`};
+      } else {
+        _data = cache[2];
+      }
+      const data = useData(_data);
+      if (c_x || c_n) {
+        return (cache[3] = <Text data={data} n={n} />);
+      } else {
+        return cache[3];
+      }
+    }
+    function useData(data) {
+      const cache = useMemoCache(2);
+      const c_data = data !== cache[0];
+      cache[0] = data;
+      let nextData;
+      if (c_data) {
+        nextData = cache[1] = {text: data.text.toLowerCase()};
+      } else {
+        nextData = cache[1];
+      }
+      return nextData;
+    }
+    let data;
+    const Text = jest.fn(function Text(props) {
+      data = props.data;
+      return data.text;
+    });
+
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      root.render(<Component />);
+    });
+    expect(root).toMatchRenderedOutput('count 0');
+    expect(Text).toBeCalledTimes(1);
+    const data0 = data;
+
+    // Changing x should reset the data object
+    await act(async () => {
+      setX(1);
+    });
+    expect(root).toMatchRenderedOutput('count 1');
+    expect(Text).toBeCalledTimes(2);
+    expect(data).not.toBe(data0);
+    const data1 = data;
+
+    // Forcing an unrelated update shouldn't recreate the
+    // data object.
+    await act(async () => {
+      forceUpdate();
+    });
+    expect(root).toMatchRenderedOutput('count 1');
+    expect(Text).toBeCalledTimes(3);
+    expect(data).toBe(data1); // confirm that the cache persisted across renders
+  });
+});