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
Joseph Savona committed
Sep 13, 2022 at 14:44 UTC
3401e9200efbdb27815f739a90f6e18d1073de13
4 files changed
+528
-12
packages/react-reconciler/src/ReactFiberHooks.new.js
+78
-6
@@ -16,7 +16,12 @@ import type {
16
Usable,
17
Thenable,
18
} from 'shared/ReactTypes';
19
-import type {Fiber, Dispatcher, HookType} from './ReactInternalTypes';
19
+import type {
20
+ Fiber,
21
+ Dispatcher,
22
+ HookType,
23
+ MemoCache,
24
+} from './ReactInternalTypes';
25
import type {Lanes, Lane} from './ReactFiberLane.new';
26
import type {HookFlags} from './ReactHookEffectTags';
27
import type {FiberRoot} from './ReactInternalTypes';
@@ -177,6 +182,8 @@ type StoreConsistencyCheck<T> = {
182
export type FunctionComponentUpdateQueue = {
183
lastEffect: Effect | null,
184
stores: Array<StoreConsistencyCheck<any>> | null,
185
+ // NOTE: optional, only set when enableUseMemoCacheHook is enabled
186
+ memoCache?: MemoCache | null,
187
};
188
189
type BasicStateAction<S> = (S => S) | S;
@@ -710,10 +717,23 @@ function updateWorkInProgressHook(): Hook {
717
return workInProgressHook;
718
}
719
713
-function createFunctionComponentUpdateQueue(): FunctionComponentUpdateQueue {
714
- return {
715
- lastEffect: null,
716
- stores: null,
720
+// NOTE: defining two versions of this function to avoid size impact when this feature is disabled.
721
+// Previously this function was inlined, the additional `memoCache` property makes it not inlined.
722
+let createFunctionComponentUpdateQueue: () => FunctionComponentUpdateQueue;
723
+if (enableUseMemoCacheHook) {
724
+ createFunctionComponentUpdateQueue = () => {
725
+ return {
726
+ lastEffect: null,
727
+ stores: null,
728
+ memoCache: null,
729
+ };
730
+ };
731
+} else {
732
+ createFunctionComponentUpdateQueue = () => {
733
+ return {
734
+ lastEffect: null,
735
+ stores: null,
736
+ };
737
};
738
}
739
@@ -787,7 +807,59 @@ function use<T>(usable: Usable<T>): T {
807
}
808
809
function useMemoCache(size: number): Array<any> {
790
- throw new Error('Not implemented.');
810
+ let memoCache = null;
811
+ // Fast-path, load memo cache from wip fiber if already prepared
812
+ let updateQueue: FunctionComponentUpdateQueue | null = (currentlyRenderingFiber.updateQueue: any);
813
+ if (updateQueue !== null) {
814
+ memoCache = updateQueue.memoCache;
815
+ }
816
+ // Otherwise clone from the current fiber
817
+ // TODO: not sure how to access the current fiber here other than going through
818
+ // currentlyRenderingFiber.alternate
819
+ if (memoCache == null) {
820
+ const current: Fiber | null = currentlyRenderingFiber.alternate;
821
+ if (current !== null) {
822
+ const currentUpdateQueue: FunctionComponentUpdateQueue | null = (current.updateQueue: any);
823
+ if (currentUpdateQueue !== null) {
824
+ const currentMemoCache: ?MemoCache = currentUpdateQueue.memoCache;
825
+ if (currentMemoCache != null) {
826
+ memoCache = {
827
+ data: currentMemoCache.data.map(array => array.slice()),
828
+ index: 0,
829
+ };
830
+ }
831
+ }
832
+ }
833
+ }
834
+ // Finally fall back to allocating a fresh instance of the cache
835
+ if (memoCache == null) {
836
+ memoCache = {
837
+ data: [],
838
+ index: 0,
839
+ };
840
+ }
841
+ if (updateQueue === null) {
842
+ updateQueue = createFunctionComponentUpdateQueue();
843
+ currentlyRenderingFiber.updateQueue = updateQueue;
844
+ }
845
+ updateQueue.memoCache = memoCache;
846
+
847
+ let data = memoCache.data[memoCache.index];
848
+ if (data === undefined) {
849
+ data = memoCache.data[memoCache.index] = new Array(size);
850
+ } else if (data.length !== size) {
851
+ // TODO: consider warning or throwing here
852
+ if (__DEV__) {
853
+ console.error(
854
+ 'Expected a constant size argument for each invocation of useMemoCache. ' +
855
+ 'The previous cache was allocated with size %s but size %s was requested.',
856
+ data.length,
857
+ size,
858
+ );
859
+ }
860
+ }
861
+ memoCache.index++;
862
+ return data;
863
}
864
865
function basicStateReducer<S>(state: S, action: BasicStateAction<S>): S {
packages/react-reconciler/src/ReactFiberHooks.old.js
+78
-6
@@ -16,7 +16,12 @@ import type {
16
Usable,
17
Thenable,
18
} from 'shared/ReactTypes';
19
-import type {Fiber, Dispatcher, HookType} from './ReactInternalTypes';
19
+import type {
20
+ Fiber,
21
+ Dispatcher,
22
+ HookType,
23
+ MemoCache,
24
+} from './ReactInternalTypes';
25
import type {Lanes, Lane} from './ReactFiberLane.old';
26
import type {HookFlags} from './ReactHookEffectTags';
27
import type {FiberRoot} from './ReactInternalTypes';
@@ -177,6 +182,8 @@ type StoreConsistencyCheck<T> = {
182
export type FunctionComponentUpdateQueue = {
183
lastEffect: Effect | null,
184
stores: Array<StoreConsistencyCheck<any>> | null,
185
+ // NOTE: optional, only set when enableUseMemoCacheHook is enabled
186
+ memoCache?: MemoCache | null,
187
};
188
189
type BasicStateAction<S> = (S => S) | S;
@@ -710,10 +717,23 @@ function updateWorkInProgressHook(): Hook {
717
return workInProgressHook;
718
}
719
713
-function createFunctionComponentUpdateQueue(): FunctionComponentUpdateQueue {
714
- return {
715
- lastEffect: null,
716
- stores: null,
720
+// NOTE: defining two versions of this function to avoid size impact when this feature is disabled.
721
+// Previously this function was inlined, the additional `memoCache` property makes it not inlined.
722
+let createFunctionComponentUpdateQueue: () => FunctionComponentUpdateQueue;
723
+if (enableUseMemoCacheHook) {
724
+ createFunctionComponentUpdateQueue = () => {
725
+ return {
726
+ lastEffect: null,
727
+ stores: null,
728
+ memoCache: null,
729
+ };
730
+ };
731
+} else {
732
+ createFunctionComponentUpdateQueue = () => {
733
+ return {
734
+ lastEffect: null,
735
+ stores: null,
736
+ };
737
};
738
}
739
@@ -787,7 +807,59 @@ function use<T>(usable: Usable<T>): T {
807
}
808
809
function useMemoCache(size: number): Array<any> {
790
- throw new Error('Not implemented.');
810
+ let memoCache = null;
811
+ // Fast-path, load memo cache from wip fiber if already prepared
812
+ let updateQueue: FunctionComponentUpdateQueue | null = (currentlyRenderingFiber.updateQueue: any);
813
+ if (updateQueue !== null) {
814
+ memoCache = updateQueue.memoCache;
815
+ }
816
+ // Otherwise clone from the current fiber
817
+ // TODO: not sure how to access the current fiber here other than going through
818
+ // currentlyRenderingFiber.alternate
819
+ if (memoCache == null) {
820
+ const current: Fiber | null = currentlyRenderingFiber.alternate;
821
+ if (current !== null) {
822
+ const currentUpdateQueue: FunctionComponentUpdateQueue | null = (current.updateQueue: any);
823
+ if (currentUpdateQueue !== null) {
824
+ const currentMemoCache: ?MemoCache = currentUpdateQueue.memoCache;
825
+ if (currentMemoCache != null) {
826
+ memoCache = {
827
+ data: currentMemoCache.data.map(array => array.slice()),
828
+ index: 0,
829
+ };
830
+ }
831
+ }
832
+ }
833
+ }
834
+ // Finally fall back to allocating a fresh instance of the cache
835
+ if (memoCache == null) {
836
+ memoCache = {
837
+ data: [],
838
+ index: 0,
839
+ };
840
+ }
841
+ if (updateQueue === null) {
842
+ updateQueue = createFunctionComponentUpdateQueue();
843
+ currentlyRenderingFiber.updateQueue = updateQueue;
844
+ }
845
+ updateQueue.memoCache = memoCache;
846
+
847
+ let data = memoCache.data[memoCache.index];
848
+ if (data === undefined) {
849
+ data = memoCache.data[memoCache.index] = new Array(size);
850
+ } else if (data.length !== size) {
851
+ // TODO: consider warning or throwing here
852
+ if (__DEV__) {
853
+ console.error(
854
+ 'Expected a constant size argument for each invocation of useMemoCache. ' +
855
+ 'The previous cache was allocated with size %s but size %s was requested.',
856
+ data.length,
857
+ size,
858
+ );
859
+ }
860
+ }
861
+ memoCache.index++;
862
+ return data;
863
}
864
865
function basicStateReducer<S>(state: S, action: BasicStateAction<S>): S {
packages/react-reconciler/src/ReactInternalTypes.js
+5
@@ -68,6 +68,11 @@ export type Dependencies = {
68
...
69
};
70
71
+export type MemoCache = {
72
+ data: Array<Array<any>>,
73
+ index: number,
74
+};
75
+
76
// A Fiber is work on a Component that needs to be done or was done. There can
77
// be more than one per component.
78
export type Fiber = {
packages/react-reconciler/src/__tests__/useMemoCache-test.js
new
+367
@@ -0,0 +1,367 @@
1
+let React;
2
+let ReactNoop;
3
+let act;
4
+let useState;
5
+let useMemoCache;
6
+let ErrorBoundary;
7
+
8
+describe('useMemoCache()', () => {
9
+ beforeEach(() => {
10
+ jest.resetModules();
11
+
12
+ React = require('react');
13
+ ReactNoop = require('react-noop-renderer');
14
+ act = require('jest-react').act;
15
+ useState = React.useState;
16
+ useMemoCache = React.unstable_useMemoCache;
17
+
18
+ class _ErrorBoundary extends React.Component {
19
+ constructor(props) {
20
+ super(props);
21
+ this.state = {hasError: false};
22
+ }
23
+
24
+ static getDerivedStateFromError(error) {
25
+ // Update state so the next render will show the fallback UI.
26
+ return {hasError: true};
27
+ }
28
+
29
+ componentDidCatch(error, errorInfo) {}
30
+
31
+ render() {
32
+ if (this.state.hasError) {
33
+ // You can render any custom fallback UI
34
+ return <h1>Something went wrong.</h1>;
35
+ }
36
+
37
+ return this.props.children;
38
+ }
39
+ }
40
+ ErrorBoundary = _ErrorBoundary;
41
+ });
42
+
43
+ // @gate enableUseMemoCacheHook
44
+ test('render component using cache', async () => {
45
+ function Component(props) {
46
+ const cache = useMemoCache(1);
47
+ expect(Array.isArray(cache)).toBe(true);
48
+ expect(cache.length).toBe(1);
49
+ expect(cache[0]).toBe(undefined);
50
+ return 'Ok';
51
+ }
52
+ const root = ReactNoop.createRoot();
53
+ await act(async () => {
54
+ root.render(<Component />);
55
+ });
56
+ expect(root).toMatchRenderedOutput('Ok');
57
+ });
58
+
59
+ // @gate enableUseMemoCacheHook
60
+ test('update component using cache', async () => {
61
+ let setX;
62
+ let forceUpdate;
63
+ function Component(props) {
64
+ const cache = useMemoCache(4);
65
+
66
+ // x is used to produce a `data` object passed to the child
67
+ const [x, _setX] = useState(0);
68
+ setX = _setX;
69
+ const c_x = x !== cache[0];
70
+ cache[0] = x;
71
+
72
+ // n is passed as-is to the child as a cache breaker
73
+ const [n, setN] = useState(0);
74
+ forceUpdate = () => setN(a => a + 1);
75
+ const c_n = n !== cache[1];
76
+ cache[1] = n;
77
+
78
+ let data;
79
+ if (c_x) {
80
+ data = cache[2] = {text: `Count ${x}`};
81
+ } else {
82
+ data = cache[2];
83
+ }
84
+ if (c_x || c_n) {
85
+ return (cache[3] = <Text data={data} n={n} />);
86
+ } else {
87
+ return cache[3];
88
+ }
89
+ }
90
+ let data;
91
+ const Text = jest.fn(function Text(props) {
92
+ data = props.data;
93
+ return data.text;
94
+ });
95
+
96
+ const root = ReactNoop.createRoot();
97
+ await act(async () => {
98
+ root.render(<Component />);
99
+ });
100
+ expect(root).toMatchRenderedOutput('Count 0');
101
+ expect(Text).toBeCalledTimes(1);
102
+ const data0 = data;
103
+
104
+ // Changing x should reset the data object
105
+ await act(async () => {
106
+ setX(1);
107
+ });
108
+ expect(root).toMatchRenderedOutput('Count 1');
109
+ expect(Text).toBeCalledTimes(2);
110
+ expect(data).not.toBe(data0);
111
+ const data1 = data;
112
+
113
+ // Forcing an unrelated update shouldn't recreate the
114
+ // data object.
115
+ await act(async () => {
116
+ forceUpdate();
117
+ });
118
+ expect(root).toMatchRenderedOutput('Count 1');
119
+ expect(Text).toBeCalledTimes(3);
120
+ expect(data).toBe(data1); // confirm that the cache persisted across renders
121
+ });
122
+
123
+ // @gate enableUseMemoCacheHook
124
+ test('update component using cache with setstate during render', async () => {
125
+ let setX;
126
+ let setN;
127
+ function Component(props) {
128
+ const cache = useMemoCache(4);
129
+
130
+ // x is used to produce a `data` object passed to the child
131
+ const [x, _setX] = useState(0);
132
+ setX = _setX;
133
+ const c_x = x !== cache[0];
134
+ cache[0] = x;
135
+
136
+ // n is passed as-is to the child as a cache breaker
137
+ const [n, _setN] = useState(0);
138
+ setN = _setN;
139
+ const c_n = n !== cache[1];
140
+ cache[1] = n;
141
+
142
+ // NOTE: setstate and early return here means that x will update
143
+ // without the data value being updated. Subsequent renders could
144
+ // therefore think that c_x = false (hasn't changed) and skip updating
145
+ // data.
146
+ // The memoizing compiler will have to handle this case, but the runtime
147
+ // can help by falling back to resetting the cache if a setstate occurs
148
+ // during render (this mirrors what we do for useMemo and friends)
149
+ if (n === 1) {
150
+ setN(2);
151
+ return;
152
+ }
153
+
154
+ let data;
155
+ if (c_x) {
156
+ data = cache[2] = {text: `Count ${x}`};
157
+ } else {
158
+ data = cache[2];
159
+ }
160
+ if (c_x || c_n) {
161
+ return (cache[3] = <Text data={data} n={n} />);
162
+ } else {
163
+ return cache[3];
164
+ }
165
+ }
166
+ let data;
167
+ const Text = jest.fn(function Text(props) {
168
+ data = props.data;
169
+ return data.text;
170
+ });
171
+
172
+ const root = ReactNoop.createRoot();
173
+ await act(async () => {
174
+ root.render(<Component />);
175
+ });
176
+ expect(root).toMatchRenderedOutput('Count 0');
177
+ expect(Text).toBeCalledTimes(1);
178
+ const data0 = data;
179
+
180
+ // Simultaneously trigger an update to x (should create a new data value)
181
+ // and trigger the setState+early return. The runtime should reset the cache
182
+ // to avoid an inconsistency
183
+ await act(async () => {
184
+ setX(1);
185
+ setN(1);
186
+ });
187
+ expect(root).toMatchRenderedOutput('Count 1');
188
+ expect(Text).toBeCalledTimes(2);
189
+ expect(data).not.toBe(data0);
190
+ const data1 = data;
191
+
192
+ // Forcing an unrelated update shouldn't recreate the
193
+ // data object.
194
+ await act(async () => {
195
+ setN(3);
196
+ });
197
+ expect(root).toMatchRenderedOutput('Count 1');
198
+ expect(Text).toBeCalledTimes(3);
199
+ expect(data).toBe(data1); // confirm that the cache persisted across renders
200
+ });
201
+
202
+ // @gate enableUseMemoCacheHook
203
+ test('update component using cache with throw during render', async () => {
204
+ let setX;
205
+ let setN;
206
+ let shouldFail = true;
207
+ function Component(props) {
208
+ const cache = useMemoCache(4);
209
+
210
+ // x is used to produce a `data` object passed to the child
211
+ const [x, _setX] = useState(0);
212
+ setX = _setX;
213
+ const c_x = x !== cache[0];
214
+ cache[0] = x;
215
+
216
+ // n is passed as-is to the child as a cache breaker
217
+ const [n, _setN] = useState(0);
218
+ setN = _setN;
219
+ const c_n = n !== cache[1];
220
+ cache[1] = n;
221
+
222
+ // NOTE the initial failure will trigger a re-render, after which the function
223
+ // will early return. This validates that the runtime resets the cache on error:
224
+ // if it doesn't the cache will be corrupt, with the cached version of data
225
+ // out of data from the cached version of x.
226
+ if (n === 1) {
227
+ if (shouldFail) {
228
+ shouldFail = false;
229
+ throw new Error('failed');
230
+ }
231
+ setN(2);
232
+ return;
233
+ }
234
+
235
+ let data;
236
+ if (c_x) {
237
+ data = cache[2] = {text: `Count ${x}`};
238
+ } else {
239
+ data = cache[2];
240
+ }
241
+ if (c_x || c_n) {
242
+ return (cache[3] = <Text data={data} n={n} />);
243
+ } else {
244
+ return cache[3];
245
+ }
246
+ }
247
+ let data;
248
+ const Text = jest.fn(function Text(props) {
249
+ data = props.data;
250
+ return data.text;
251
+ });
252
+
253
+ spyOnDev(console, 'error');
254
+
255
+ const root = ReactNoop.createRoot();
256
+ await act(async () => {
257
+ root.render(
258
+ <ErrorBoundary>
259
+ <Component />
260
+ </ErrorBoundary>,
261
+ );
262
+ });
263
+ expect(root).toMatchRenderedOutput('Count 0');
264
+ expect(Text).toBeCalledTimes(1);
265
+ const data0 = data;
266
+
267
+ // Simultaneously trigger an update to x (should create a new data value)
268
+ // and trigger the setState+early return. The runtime should reset the cache
269
+ // to avoid an inconsistency
270
+ await act(async () => {
271
+ // this update bumps the count
272
+ setX(1);
273
+ // this triggers a throw.
274
+ setN(1);
275
+ });
276
+ expect(root).toMatchRenderedOutput('Count 1');
277
+ expect(Text).toBeCalledTimes(2);
278
+ expect(data).not.toBe(data0);
279
+ const data1 = data;
280
+
281
+ // Forcing an unrelated update shouldn't recreate the
282
+ // data object.
283
+ await act(async () => {
284
+ setN(3);
285
+ });
286
+ expect(root).toMatchRenderedOutput('Count 1');
287
+ expect(Text).toBeCalledTimes(3);
288
+ expect(data).toBe(data1); // confirm that the cache persisted across renders
289
+ });
290
+
291
+ // @gate enableUseMemoCacheHook
292
+ test('update component and custom hook with caches', async () => {
293
+ let setX;
294
+ let forceUpdate;
295
+ function Component(props) {
296
+ const cache = useMemoCache(4);
297
+
298
+ // x is used to produce a `data` object passed to the child
299
+ const [x, _setX] = useState(0);
300
+ setX = _setX;
301
+ const c_x = x !== cache[0];
302
+ cache[0] = x;
303
+
304
+ // n is passed as-is to the child as a cache breaker
305
+ const [n, setN] = useState(0);
306
+ forceUpdate = () => setN(a => a + 1);
307
+ const c_n = n !== cache[1];
308
+ cache[1] = n;
309
+
310
+ let _data;
311
+ if (c_x) {
312
+ _data = cache[2] = {text: `Count ${x}`};
313
+ } else {
314
+ _data = cache[2];
315
+ }
316
+ const data = useData(_data);
317
+ if (c_x || c_n) {
318
+ return (cache[3] = <Text data={data} n={n} />);
319
+ } else {
320
+ return cache[3];
321
+ }
322
+ }
323
+ function useData(data) {
324
+ const cache = useMemoCache(2);
325
+ const c_data = data !== cache[0];
326
+ cache[0] = data;
327
+ let nextData;
328
+ if (c_data) {
329
+ nextData = cache[1] = {text: data.text.toLowerCase()};
330
+ } else {
331
+ nextData = cache[1];
332
+ }
333
+ return nextData;
334
+ }
335
+ let data;
336
+ const Text = jest.fn(function Text(props) {
337
+ data = props.data;
338
+ return data.text;
339
+ });
340
+
341
+ const root = ReactNoop.createRoot();
342
+ await act(async () => {
343
+ root.render(<Component />);
344
+ });
345
+ expect(root).toMatchRenderedOutput('count 0');
346
+ expect(Text).toBeCalledTimes(1);
347
+ const data0 = data;
348
+
349
+ // Changing x should reset the data object
350
+ await act(async () => {
351
+ setX(1);
352
+ });
353
+ expect(root).toMatchRenderedOutput('count 1');
354
+ expect(Text).toBeCalledTimes(2);
355
+ expect(data).not.toBe(data0);
356
+ const data1 = data;
357
+
358
+ // Forcing an unrelated update shouldn't recreate the
359
+ // data object.
360
+ await act(async () => {
361
+ forceUpdate();
362
+ });
363
+ expect(root).toMatchRenderedOutput('count 1');
364
+ expect(Text).toBeCalledTimes(3);
365
+ expect(data).toBe(data1); // confirm that the cache persisted across renders
366
+ });
367
+});