main
js 141 lines 4.4 KB
Raw
1 /**
2 * Copyright (c) Meta Platforms, Inc. and its affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 *
7 * @flow
8 */
9
10 import ReactSharedInternals from 'shared/ReactSharedInternals';
11
12 const UNTERMINATED = 0;
13 const TERMINATED = 1;
14 const ERRORED = 2;
15
16 type UnterminatedCacheNode<T> = {
17 s: 0,
18 v: void,
19 o: null | WeakMap<Function | Object, CacheNode<T>>,
20 p: null | Map<string | number | null | void | symbol | boolean, CacheNode<T>>,
21 };
22
23 type TerminatedCacheNode<T> = {
24 s: 1,
25 v: T,
26 o: null | WeakMap<Function | Object, CacheNode<T>>,
27 p: null | Map<string | number | null | void | symbol | boolean, CacheNode<T>>,
28 };
29
30 type ErroredCacheNode<T> = {
31 s: 2,
32 v: mixed,
33 o: null | WeakMap<Function | Object, CacheNode<T>>,
34 p: null | Map<string | number | null | void | symbol | boolean, CacheNode<T>>,
35 };
36
37 type CacheNode<T> =
38 | TerminatedCacheNode<T>
39 | UnterminatedCacheNode<T>
40 | ErroredCacheNode<T>;
41
42 function createCacheRoot<T>(): WeakMap<Function | Object, CacheNode<T>> {
43 return new WeakMap();
44 }
45
46 function createCacheNode<T>(): CacheNode<T> {
47 return {
48 s: UNTERMINATED, // status, represents whether the cached computation returned a value or threw an error
49 v: undefined, // value, either the cached result or an error, depending on s
50 o: null, // object cache, a WeakMap where non-primitive arguments are stored
51 p: null, // primitive cache, a regular Map where primitive arguments are stored.
52 };
53 }
54
55 export function cache<A: Iterable<mixed>, T>(fn: (...A) => T): (...A) => T {
56 return function () {
57 const dispatcher = ReactSharedInternals.A;
58 if (!dispatcher) {
59 // If there is no dispatcher, then we treat this as not being cached.
60 // $FlowFixMe[incompatible-type]: We don't want to use rest arguments since we transpile the code.
61 return fn.apply(null, arguments);
62 }
63 const fnMap: WeakMap<any, CacheNode<T>> = dispatcher.getCacheForType(
64 createCacheRoot,
65 );
66 const fnNode = fnMap.get(fn);
67 let cacheNode: CacheNode<T>;
68 if (fnNode === undefined) {
69 cacheNode = createCacheNode();
70 fnMap.set(fn, cacheNode);
71 } else {
72 cacheNode = fnNode;
73 }
74 for (let i = 0, l = arguments.length; i < l; i++) {
75 const arg = arguments[i];
76 if (
77 typeof arg === 'function' ||
78 // $FlowFixMe[invalid-compare]
79 (typeof arg === 'object' && arg !== null)
80 ) {
81 // Objects go into a WeakMap
82 let objectCache = cacheNode.o;
83 if (objectCache === null) {
84 cacheNode.o = objectCache = new WeakMap();
85 }
86 const objectNode = objectCache.get(arg);
87 if (objectNode === undefined) {
88 cacheNode = createCacheNode();
89 objectCache.set(arg, cacheNode);
90 } else {
91 cacheNode = objectNode;
92 }
93 } else {
94 // Primitives go into a regular Map
95 let primitiveCache = cacheNode.p;
96 if (primitiveCache === null) {
97 cacheNode.p = primitiveCache = new Map();
98 }
99 const primitiveNode = primitiveCache.get(arg);
100 if (primitiveNode === undefined) {
101 cacheNode = createCacheNode();
102 primitiveCache.set(arg, cacheNode);
103 } else {
104 cacheNode = primitiveNode;
105 }
106 }
107 }
108 if (cacheNode.s === TERMINATED) {
109 return cacheNode.v;
110 }
111 if (cacheNode.s === ERRORED) {
112 throw cacheNode.v;
113 }
114 try {
115 // $FlowFixMe[incompatible-type]: We don't want to use rest arguments since we transpile the code.
116 const result = fn.apply(null, arguments);
117 const terminatedNode: TerminatedCacheNode<T> = cacheNode as any;
118 terminatedNode.s = TERMINATED;
119 terminatedNode.v = result;
120 return result;
121 } catch (error) {
122 // We store the first error that's thrown and rethrow it.
123 const erroredNode: ErroredCacheNode<T> = cacheNode as any;
124 erroredNode.s = ERRORED;
125 erroredNode.v = error;
126 throw error;
127 }
128 };
129 }
130
131 export function cacheSignal(): null | AbortSignal {
132 const dispatcher = ReactSharedInternals.A;
133 if (!dispatcher) {
134 // If there is no dispatcher, then we treat this as not having an AbortSignal
135 // since in the same context, a cached function will be allowed to be called
136 // but it won't be cached. So it's neither an infinite AbortSignal nor an
137 // already resolved one.
138 return null;
139 }
140 return dispatcher.cacheSignal();
141 }