@samitouri / QOS-React-1 / commits / 6b23c25ff9

Support HMR

Adds supports for hot module reloading (HMR) by resetting the cache if a hash of the source file changes. This is enabled via a compiler flag, but also enabled automatically via the babel plugin when NODE_ENV=development. ghstack-source-id: 5cd1ad5c893533c93275c2d78b57ec0c8f1d24e7 Pull Request resolved: https://github.com/facebook/react-forget/pull/2951

Joe Savona committed May 9, 2024 at 13:48 UTC 6b23c25ff93dafce938aa94910b9a94df33bc850
9 files changed +225 -8
compiler/apps/playground/components/Editor/EditorImpl.tsx
+1
@@ -204,6 +204,7 @@ function compile(source: string): CompilerOutput {
204 getReactFunctionType(id),
205 null,
206 null,
207 + null,
208 )) {
209 const fnName = fn.node.id?.name ?? null;
210 switch (result.kind) {
compiler/packages/babel-plugin-react-compiler/src/Babel/BabelPlugin.ts
+10
@@ -33,10 +33,20 @@ export default function BabelPluginReactCompiler(
33 if (pipelineUsesReanimatedPlugin(pass.file.opts.plugins)) {
34 opts = injectReanimatedFlag(opts);
35 }
36 + if (process.env["NODE_ENV"] === "development") {
37 + opts = {
38 + ...opts,
39 + environment: {
40 + ...opts.environment,
41 + enableResetCacheOnSourceFileChanges: true,
42 + },
43 + };
44 + }
45 compileProgram(prog, {
46 opts,
47 filename: pass.filename ?? null,
48 comments: pass.file.ast.comments ?? [],
49 + code: pass.file.code,
50 });
51 },
52 },
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
+7 -4
@@ -106,7 +106,8 @@ export function* run(
106 fnType: ReactFunctionType,
107
108 logger: Logger | null,
109 - filename: string | null
109 + filename: string | null,
110 + code: string | null
111 ): Generator<CompilerPipelineValue, CodegenFunction> {
112 const contextIdentifiers = findContextIdentifiers(func);
113 const env = new Environment(
@@ -114,7 +115,8 @@ export function* run(
115 config,
116 contextIdentifiers,
117 logger,
117 - filename
118 + filename,
119 + code
120 );
121 yield {
122 kind: "debug",
@@ -452,9 +454,10 @@ export function compileFn(
454 config: EnvironmentConfig,
455 fnType: ReactFunctionType,
456 logger: Logger | null,
455 - filename: string | null
457 + filename: string | null,
458 + code: string | null
459 ): CodegenFunction {
457 - let generator = run(func, config, fnType, logger, filename);
460 + let generator = run(func, config, fnType, logger, filename, code);
461 while (true) {
462 const next = generator.next();
463 if (next.done) {
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
+9 -1
@@ -36,6 +36,7 @@ export type CompilerPass = {
36 opts: PluginOptions;
37 filename: string | null;
38 comments: Array<t.CommentBlock | t.CommentLine>;
39 + code: string | null;
40 };
41
42 function findDirectiveEnablingMemoization(
@@ -308,7 +309,14 @@ export function compileProgram(
309 }
310 const config = environment.unwrap();
311
311 - compiledFn = compileFn(fn, config, fnType, options.logger, pass.filename);
312 + compiledFn = compileFn(
313 + fn,
314 + config,
315 + fnType,
316 + options.logger,
317 + pass.filename,
318 + pass.code
319 + );
320 options.logger?.logEvent(pass.filename, {
321 kind: "CompileSuccess",
322 fnLoc: fn.node.loc ?? null,
compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
+11 -1
@@ -119,6 +119,13 @@ export type Hook = z.infer<typeof HookSchema>;
119 const EnvironmentConfigSchema = z.object({
120 customHooks: z.map(z.string(), HookSchema).optional().default(new Map()),
121
122 + /**
123 + * Enable a check that resets the memoization cache when the source code of the file changes.
124 + * This is intended to support hot module reloading (HMR), where the same runtime component
125 + * instance will be reused across different versions of the component source.
126 + */
127 + enableResetCacheOnSourceFileChanges: z.boolean().default(false),
128 +
129 /**
130 * Enable using information from existing useMemo/useCallback to understand when a value is done
131 * being mutated. With this mode enabled, Forget will still discard the actual useMemo/useCallback
@@ -427,6 +434,7 @@ export class Environment {
434 #nextScope: number = 0;
435 logger: Logger | null;
436 filename: string | null;
437 + code: string | null;
438 config: EnvironmentConfig;
439 fnType: ReactFunctionType;
440
@@ -438,11 +446,13 @@ export class Environment {
446 config: EnvironmentConfig,
447 contextIdentifiers: Set<t.Identifier>,
448 logger: Logger | null,
441 - filename: string | null
449 + filename: string | null,
450 + code: string | null
451 ) {
452 this.fnType = fnType;
453 this.config = config;
454 this.filename = filename;
455 + this.code = code;
456 this.logger = logger;
457 this.#shapes = new Map(DEFAULT_SHAPES);
458 this.#globals = new Map(DEFAULT_GLOBALS);
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts
+88 -1
@@ -43,6 +43,7 @@ import { assertExhaustive } from "../Utils/utils";
43 import { buildReactiveFunction } from "./BuildReactiveFunction";
44 import { SINGLE_CHILD_FBT_TAGS } from "./MemoizeFbtOperandsInSameScope";
45 import { ReactiveFunctionVisitor, visitReactiveFunction } from "./visitors";
46 +import { createHmac } from "crypto";
47
48 export const MEMO_CACHE_SENTINEL = "react.memo_cache_sentinel";
49 export const EARLY_RETURN_SENTINEL = "react.early_return_sentinel";
@@ -78,6 +79,25 @@ export function codegenFunction(
79 uniqueIdentifiers,
80 null
81 );
82 +
83 + /**
84 + * Hot-module reloading reuses component instances at runtime even as the source of the component changes.
85 + * The generated code needs to prevent values from one version of the code being reused after a code cange.
86 + * If HMR detection is enabled and we know the source code of the component, assign a cache slot to track
87 + * the source hash, and later, emit code to check for source changes and reset the cache on source changes.
88 + */
89 + let hotModuleReloadState: { cacheIndex: number; hash: string } | null = null;
90 + if (
91 + fn.env.config.enableResetCacheOnSourceFileChanges &&
92 + fn.env.code !== null
93 + ) {
94 + const hash = createHmac("sha256", fn.env.code).digest("hex");
95 + hotModuleReloadState = {
96 + cacheIndex: cx.nextCacheIndex,
97 + hash,
98 + };
99 + }
100 +
101 const compileResult = codegenReactiveFunction(cx, fn);
102 if (compileResult.isErr()) {
103 return compileResult;
@@ -98,8 +118,10 @@ export function codegenFunction(
118
119 const cacheCount = compiled.memoSlotsUsed;
120 if (cacheCount !== 0) {
121 + const preface: Array<t.Statement> = [];
122 +
123 // The import declaration for `useMemoCache` is inserted in the Babel plugin
102 - compiled.body.body.unshift(
124 + preface.push(
125 t.variableDeclaration("const", [
126 t.variableDeclarator(
127 t.identifier(cx.synthesizeName("$")),
@@ -109,6 +131,71 @@ export function codegenFunction(
131 ),
132 ])
133 );
134 + if (hotModuleReloadState !== null) {
135 + // HMR detection is enabled, emit code to reset the memo cache on source changes
136 + const index = cx.synthesizeName("$i");
137 + preface.push(
138 + t.ifStatement(
139 + t.binaryExpression(
140 + "!==",
141 + t.memberExpression(
142 + t.identifier(cx.synthesizeName("$")),
143 + t.numericLiteral(hotModuleReloadState.cacheIndex),
144 + true
145 + ),
146 + t.stringLiteral(hotModuleReloadState.hash)
147 + ),
148 + t.blockStatement([
149 + t.forStatement(
150 + t.variableDeclaration("let", [
151 + t.variableDeclarator(t.identifier(index), t.numericLiteral(0)),
152 + ]),
153 + t.binaryExpression(
154 + "<",
155 + t.identifier(index),
156 + t.numericLiteral(cacheCount)
157 + ),
158 + t.assignmentExpression(
159 + "+=",
160 + t.identifier(index),
161 + t.numericLiteral(1)
162 + ),
163 + t.blockStatement([
164 + t.expressionStatement(
165 + t.assignmentExpression(
166 + "=",
167 + t.memberExpression(
168 + t.identifier(cx.synthesizeName("$")),
169 + t.identifier(index),
170 + true
171 + ),
172 + t.callExpression(
173 + t.memberExpression(
174 + t.identifier("Symbol"),
175 + t.identifier("for")
176 + ),
177 + [t.stringLiteral(MEMO_CACHE_SENTINEL)]
178 + )
179 + )
180 + ),
181 + ])
182 + ),
183 + t.expressionStatement(
184 + t.assignmentExpression(
185 + "=",
186 + t.memberExpression(
187 + t.identifier(cx.synthesizeName("$")),
188 + t.numericLiteral(hotModuleReloadState.cacheIndex),
189 + true
190 + ),
191 + t.stringLiteral(hotModuleReloadState.hash)
192 + )
193 + ),
194 + ])
195 + )
196 + );
197 + }
198 + compiled.body.body.unshift(...preface);
199 }
200
201 const emitInstrumentForget = fn.env.config.enableEmitInstrumentForget;
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hot-module-reloading.expect.md new
+83
@@ -0,0 +1,83 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @enableResetCacheOnSourceFileChanges
6 +import { useMemo, useState } from "react";
7 +import { ValidateMemoization } from "shared-runtime";
8 +
9 +function Component(props) {
10 + const [state, setState] = useState(0);
11 + const doubled = useMemo(() => [state * 2], [state]);
12 + return <ValidateMemoization inputs={[state]} output={doubled} />;
13 +}
14 +
15 +export const FIXTURE_ENTRYPOINT = {
16 + fn: Component,
17 + params: [{}],
18 + sequentialRenders: [{}, {}],
19 +};
20 +
21 +```
22 +
23 +## Code
24 +
25 +```javascript
26 +import { c as useMemoCache } from "react/compiler-runtime"; // @enableResetCacheOnSourceFileChanges
27 +import { useMemo, useState } from "react";
28 +import { ValidateMemoization } from "shared-runtime";
29 +
30 +function Component(props) {
31 + const $ = useMemoCache(8);
32 + if (
33 + $[0] !== "bb6936608c0afe8e313aa547ca09fbc8451f24664284368812127c7e9bc2bca9"
34 + ) {
35 + for (let $i = 0; $i < 8; $i += 1) {
36 + $[$i] = Symbol.for("react.memo_cache_sentinel");
37 + }
38 + $[0] = "bb6936608c0afe8e313aa547ca09fbc8451f24664284368812127c7e9bc2bca9";
39 + }
40 + const [state] = useState(0);
41 + let t0;
42 + const t1 = state * 2;
43 + let t2;
44 + if ($[1] !== t1) {
45 + t2 = [t1];
46 + $[1] = t1;
47 + $[2] = t2;
48 + } else {
49 + t2 = $[2];
50 + }
51 + t0 = t2;
52 + const doubled = t0;
53 + let t3;
54 + if ($[3] !== state) {
55 + t3 = [state];
56 + $[3] = state;
57 + $[4] = t3;
58 + } else {
59 + t3 = $[4];
60 + }
61 + let t4;
62 + if ($[5] !== t3 || $[6] !== doubled) {
63 + t4 = <ValidateMemoization inputs={t3} output={doubled} />;
64 + $[5] = t3;
65 + $[6] = doubled;
66 + $[7] = t4;
67 + } else {
68 + t4 = $[7];
69 + }
70 + return t4;
71 +}
72 +
73 +export const FIXTURE_ENTRYPOINT = {
74 + fn: Component,
75 + params: [{}],
76 + sequentialRenders: [{}, {}],
77 +};
78 +
79 +```
80 +
81 +### Eval output
82 +(kind: ok) <div>{"inputs":[0],"output":[0]}</div>
83 +<div>{"inputs":[0],"output":[0]}</div>
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hot-module-reloading.js new
+15
@@ -0,0 +1,15 @@
1 +// @enableResetCacheOnSourceFileChanges
2 +import { useMemo, useState } from "react";
3 +import { ValidateMemoization } from "shared-runtime";
4 +
5 +function Component(props) {
6 + const [state, setState] = useState(0);
7 + const doubled = useMemo(() => [state * 2], [state]);
8 + return <ValidateMemoization inputs={[state]} output={doubled} />;
9 +}
10 +
11 +export const FIXTURE_ENTRYPOINT = {
12 + fn: Component,
13 + params: [{}],
14 + sequentialRenders: [{}, {}],
15 +};
compiler/packages/snap/src/fixture-utils.ts
+1 -1
@@ -37,7 +37,7 @@ function stripExtension(filename: string, extensions: Array<string>): string {
37
38 export async function readTestFilter(): Promise<TestFilter | null> {
39 if (!(await exists(FILTER_PATH))) {
40 - throw new Error(`testfilter file not found at ${FILTER_PATH}`);
40 + throw new Error(`testfilter file not found at \`${FILTER_PATH}\``);
41 }
42
43 const input = await fs.readFile(FILTER_PATH, "utf8");