From f71713fb1b08d14750a164549775328e34d45cb7 Mon Sep 17 00:00:00 2001 From: Joseph Savona Date: Tue, 13 Sep 2022 14:44:32 -0700 Subject: [PATCH] useMemoCache implementation (#25143) * 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 --- .../src/ReactFiberHooks.new.js | 84 +++- .../src/ReactFiberHooks.old.js | 84 +++- .../src/ReactInternalTypes.js | 5 + .../src/__tests__/useMemoCache-test.js | 367 ++++++++++++++++++ 4 files changed, 528 insertions(+), 12 deletions(-) create mode 100644 packages/react-reconciler/src/__tests__/useMemoCache-test.js diff --git a/packages/react-reconciler/src/ReactFiberHooks.new.js b/packages/react-reconciler/src/ReactFiberHooks.new.js index 3f45ec61434a9..602f855c76d85 100644 --- 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 = { export type FunctionComponentUpdateQueue = { lastEffect: Effect | null, stores: Array> | null, + // NOTE: optional, only set when enableUseMemoCacheHook is enabled + memoCache?: MemoCache | null, }; type BasicStateAction = (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(usable: Usable): T { } function useMemoCache(size: number): Array { - 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(state: S, action: BasicStateAction): S { diff --git a/packages/react-reconciler/src/ReactFiberHooks.old.js b/packages/react-reconciler/src/ReactFiberHooks.old.js index 5495a2992b787..f2b024bee943d 100644 --- 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 = { export type FunctionComponentUpdateQueue = { lastEffect: Effect | null, stores: Array> | null, + // NOTE: optional, only set when enableUseMemoCacheHook is enabled + memoCache?: MemoCache | null, }; type BasicStateAction = (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(usable: Usable): T { } function useMemoCache(size: number): Array { - 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(state: S, action: BasicStateAction): S { diff --git a/packages/react-reconciler/src/ReactInternalTypes.js b/packages/react-reconciler/src/ReactInternalTypes.js index df7398c7963d0..004aa17e56ad8 100644 --- 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>, + 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 new file mode 100644 index 0000000000000..94684ec8c00b3 --- /dev/null +++ b/packages/react-reconciler/src/__tests__/useMemoCache-test.js @@ -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

Something went wrong.

; + } + + 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(); + }); + 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] = ); + } 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(); + }); + 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] = ); + } 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(); + }); + 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] = ); + } 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( + + + , + ); + }); + 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] = ); + } 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(); + }); + 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 + }); +});