@samitouri / QOS-React-1 / commits / ac170f4f0d

[babel] Allow userspace useMemoCache

Not every consumer of Forget will be able to run an experimental version of React. In the meantime before useMemoCache is stable, provide a way for OSS to pass in a userspace impl.

Lauren Tan committed Oct 31, 2023 at 12:16 UTC ac170f4f0d8790ee607dbf8bd33ae76778a93ac7
7 files changed +157 -10
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Imports.ts
+18
@@ -145,3 +145,21 @@ export function insertUseMemoCacheImportDeclaration(
145 )
146 );
147 }
148 +
149 +export function insertUserspaceUseMemoCacheImportDeclaration(
150 + program: NodePath<t.Program>,
151 + moduleName: string
152 +): void {
153 + program.unshiftContainer(
154 + "body",
155 + t.importDeclaration(
156 + [
157 + t.importSpecifier(
158 + t.identifier("useMemoCache"),
159 + t.identifier("unstable_useMemoCache")
160 + ),
161 + ],
162 + t.stringLiteral(moduleName)
163 + )
164 + );
165 +}
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Options.ts
+12
@@ -109,6 +109,17 @@ export type PluginOptions = {
109 * ```
110 */
111 compilationMode: CompilationMode;
112 +
113 + /**
114 + * If specified, Forget will import `useMemoCache` from this module instead of React. Use this if
115 + * you are for whatever reason unable to use an experimental version of React.
116 + *
117 + * ```
118 + * // If specified:
119 + * import {unstable_useMemoCache} from 'useMemoCache_DO_NOT_USE';
120 + * ```
121 + */
122 + useMemoCacheSource: string | null;
123 };
124
125 export type CompilationMode =
@@ -156,6 +167,7 @@ export const defaultOptions: PluginOptions = {
167 gating: null,
168 instrumentForget: null,
169 noEmit: false,
170 + useMemoCacheSource: null,
171 } as const;
172
173 export function parsePluginOptions(obj: unknown): PluginOptions {
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
+18 -10
@@ -21,10 +21,11 @@ import {
21 addImportsToProgram,
22 findExistingImports,
23 insertUseMemoCacheImportDeclaration,
24 + insertUserspaceUseMemoCacheImportDeclaration,
25 updateExistingReactImportDeclaration,
26 } from "./Imports";
27 import { addInstrumentForget } from "./Instrumentation";
27 -import { PluginOptions, parsePluginOptions } from "./Options";
28 +import { ExternalFunction, PluginOptions, parsePluginOptions } from "./Options";
29 import { compileFn } from "./Pipeline";
30
31 export type CompilerPass = {
@@ -325,18 +326,25 @@ export function compileProgram(
326 // `import {unstable_useMemoCache as useMemoCache} from 'react'` and rename
327 // `React.unstable_useMemoCache(n)` to `useMemoCache(n)`;
328 if (didInsertUseMemoCache) {
328 - if (hasExistingReactImport) {
329 - const didUpdateImport = updateExistingReactImportDeclaration(program);
330 - if (didUpdateImport === false) {
331 - throw new Error(
332 - "Expected an ImportDeclaration of react in order to update ImportSpecifiers with useMemoCache"
333 - );
329 + if (options.useMemoCacheSource === null) {
330 + if (hasExistingReactImport) {
331 + const didUpdateImport = updateExistingReactImportDeclaration(program);
332 + if (didUpdateImport === false) {
333 + throw new Error(
334 + "Expected an ImportDeclaration of react in order to update ImportSpecifiers with useMemoCache"
335 + );
336 + }
337 + } else {
338 + insertUseMemoCacheImportDeclaration(program);
339 }
335 - } else {
336 - insertUseMemoCacheImportDeclaration(program);
340 + } else if (typeof options.useMemoCacheSource === "string") {
341 + insertUserspaceUseMemoCacheImportDeclaration(
342 + program,
343 + options.useMemoCacheSource
344 + );
345 }
346 }
339 - const externalFunctions = [];
347 + const externalFunctions: ExternalFunction[] = [];
348 // TODO: check for duplicate import specifiers
349 if (options.gating != null) {
350 externalFunctions.push(options.gating);
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/userspace-use-memo-cache.expect.md new
+73
@@ -0,0 +1,73 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @useMemoCacheSource
6 +function Component(props) {
7 + const [x, setX] = useState(1);
8 + let y;
9 + if (props.cond) {
10 + y = x * 2;
11 + }
12 + return (
13 + <Button
14 + onClick={() => {
15 + setX(10 * y);
16 + }}
17 + ></Button>
18 + );
19 +}
20 +
21 +export const FIXTURE_ENTRYPOINT = {
22 + fn: Component,
23 + params: [true],
24 + isComponent: true,
25 +};
26 +
27 +```
28 +
29 +## Code
30 +
31 +```javascript
32 +import { unstable_useMemoCache as useMemoCache } from "shared-runtime"; // @useMemoCacheSource
33 +function Component(props) {
34 + const $ = useMemoCache(5);
35 + const [x, setX] = useState(1);
36 + let y;
37 + if ($[0] !== props.cond || $[1] !== x) {
38 + if (props.cond) {
39 + y = x * 2;
40 + }
41 + $[0] = props.cond;
42 + $[1] = x;
43 + $[2] = y;
44 + } else {
45 + y = $[2];
46 + }
47 +
48 + const t0 = y;
49 + let t1;
50 + if ($[3] !== t0) {
51 + t1 = (
52 + <Button
53 + onClick={() => {
54 + setX(10 * y);
55 + }}
56 + />
57 + );
58 + $[3] = t0;
59 + $[4] = t1;
60 + } else {
61 + t1 = $[4];
62 + }
63 + return t1;
64 +}
65 +
66 +export const FIXTURE_ENTRYPOINT = {
67 + fn: Component,
68 + params: [true],
69 + isComponent: true,
70 +};
71 +
72 +```
73 +
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/userspace-use-memo-cache.js new
+21
@@ -0,0 +1,21 @@
1 +// @useMemoCacheSource
2 +function Component(props) {
3 + const [x, setX] = useState(1);
4 + let y;
5 + if (props.cond) {
6 + y = x * 2;
7 + }
8 + return (
9 + <Button
10 + onClick={() => {
11 + setX(10 * y);
12 + }}
13 + ></Button>
14 + );
15 +}
16 +
17 +export const FIXTURE_ENTRYPOINT = {
18 + fn: Component,
19 + params: [true],
20 + isComponent: true,
21 +};
compiler/packages/fixture-test-utils/src/compiler-utils.ts
+5
@@ -24,6 +24,7 @@ export function transformFixtureInput(
24 let instrumentForget = null;
25 let enableEmitFreeze = null;
26 let compilationMode: CompilationMode = "all";
27 + let useMemoCacheSource = null;
28
29 if (firstLine.indexOf("@compilationMode(annotation)") !== -1) {
30 assert(
@@ -58,6 +59,9 @@ export function transformFixtureInput(
59 importSpecifierName: "makeReadOnly",
60 };
61 }
62 + if (firstLine.includes("@useMemoCacheSource")) {
63 + useMemoCacheSource = "shared-runtime";
64 + }
65 const config = parseConfigPragmaFn(firstLine);
66 const result = pluginFn(
67 input,
@@ -104,6 +108,7 @@ export function transformFixtureInput(
108 instrumentForget,
109 panicThreshold: "ALL_ERRORS",
110 noEmit: false,
111 + useMemoCacheSource,
112 },
113 includeAst
114 );
compiler/packages/sprout/src/shared-runtime.ts
+10
@@ -191,3 +191,13 @@ export const ObjectWithHooks = {
191 return 0;
192 },
193 };
194 +
195 +const $empty = Symbol.for("react.memo_cache_sentinel");
196 +export function unstable_useMemoCache(size: number) {
197 + "use no forget";
198 + const $ = new Array(size);
199 + for (let ii = 0; ii < size; ii++) {
200 + $[ii] = $empty;
201 + }
202 + return React.useRef($).current;
203 +}