From 4ce22ebcfeb264b6b20e5a5b6f7fc16585a0f907 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Fri, 19 May 2023 11:05:04 -0700 Subject: [PATCH] Option to assume hooks follow the rules MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a new feature flag which tells the compiler to assume that hooks follow the Rules of React. Specifically, the idea that since any hook could be wrapped in a giant `useMemo()` call, all arguments to hooks have to be treated as if they're owned by React — and therefore become immutable — and that the return value of the hook is immutable. Our default is to assume that hooks break the rules, but in practice nearly every component follows them. --- .../packages/snap/src/compiler-worker.ts | 5 ++ compiler/forget/src/HIR/Environment.ts | 52 ++++++++++++++++--- compiler/forget/src/HIR/Types.ts | 2 +- .../forget/src/Inference/DropMemoCalls.ts | 2 +- .../src/Inference/InferReferenceEffects.ts | 19 ++++--- .../forget/src/__tests__/compiler-test.ts | 3 ++ .../compiler/immutable-hooks.expect.md | 48 +++++++++++++++++ .../fixtures/compiler/immutable-hooks.js | 9 ++++ .../test-utils/generateTestsFromFixtures.ts | 13 ++++- 9 files changed, 135 insertions(+), 18 deletions(-) create mode 100644 compiler/forget/src/__tests__/fixtures/compiler/immutable-hooks.expect.md create mode 100644 compiler/forget/src/__tests__/fixtures/compiler/immutable-hooks.js diff --git a/compiler/forget/packages/snap/src/compiler-worker.ts b/compiler/forget/packages/snap/src/compiler-worker.ts index 9d32e9eb87..7f3bb45bf8 100644 --- a/compiler/forget/packages/snap/src/compiler-worker.ts +++ b/compiler/forget/packages/snap/src/compiler-worker.ts @@ -93,6 +93,7 @@ export async function compile( let instrumentForget = null; let panicOnBailout = true; let memoizeJsxElements = true; + let enableAssumeHooksFollowRulesOfReact = false; if (firstLine.indexOf("@forgetDirective") !== -1) { enableOnlyOnUseForgetDirective = true; } @@ -120,6 +121,9 @@ export async function compile( if (firstLine.indexOf("@memoizeJsxElements false") !== -1) { memoizeJsxElements = false; } + if (firstLine.indexOf("@enableAssumeHooksFollowRulesOfReact true") !== -1) { + enableAssumeHooksFollowRulesOfReact = true; + } const language = parseLanguage(firstLine); @@ -139,6 +143,7 @@ export async function compile( ]), validateHooksUsage: true, enableFunctionCallSignatureOptimizations: true, + enableAssumeHooksFollowRulesOfReact, inlineUseMemo: true, memoizeJsxElements, }, diff --git a/compiler/forget/src/HIR/Environment.ts b/compiler/forget/src/HIR/Environment.ts index f942fc0dab..acf1343de4 100644 --- a/compiler/forget/src/HIR/Environment.ts +++ b/compiler/forget/src/HIR/Environment.ts @@ -17,12 +17,10 @@ import { import { BlockId, BuiltInType, - Effect, FunctionType, IdentifierId, ObjectType, PolyType, - ValueKind, makeBlockId, makeIdentifierId, } from "./HIR"; @@ -38,10 +36,50 @@ import { FunctionSignature, ShapeRegistry } from "./ObjectShape"; // missing some recursive Object / Function shapeIds export type EnvironmentConfig = Partial<{ customHooks: Map; + + /** + * Enable memoization of JSX elements in addition to other types of values. When disabled, + * other types (objects, arrays, call expressions, etc) are memoized, but not known JSX + * values. + * + * Defaults to true + */ memoizeJsxElements: boolean; + + /** + * Enable validation of hooks to partially check that the component honors the rules of hooks. + * When disabled, the component is assumed to follow the rules (though the Babel plugin looks + * for suppressions of the lint rule). + * + * Defaults to false + */ validateHooksUsage: boolean; + + /** + * Enable inlining of `useMemo()` function expressions so that they can be more optimally + * compiled. + * + * Defaults to false + */ inlineUseMemo: boolean; + + /** + * Enable optimizations based on the signature of (non-method) built-in function calls. + * + * Defaults to false + */ enableFunctionCallSignatureOptimizations: boolean; + + /** + * When enabled, the compiler assumes that hooks follow the Rules of React: + * - Hooks may memoize computation based on any of their parameters, thus + * any arguments to a hook are assumed frozen after calling the hook. + * - Hooks may memoize the result they return, thus the return value is + * assumed frozen. + + * Defaults to false + */ + enableAssumeHooksFollowRulesOfReact: boolean; }>; export class Environment { @@ -51,6 +89,7 @@ export class Environment { #nextBlock: number = 0; validateHooksUsage: boolean; enableFunctionCallSignatureOptimizations: boolean; + enableAssumeHooksFollowRulesOfReact: boolean; #contextIdentifiers: Set; constructor( @@ -77,6 +116,8 @@ export class Environment { this.validateHooksUsage = config?.validateHooksUsage ?? false; this.enableFunctionCallSignatureOptimizations = config?.enableFunctionCallSignatureOptimizations ?? false; + this.enableAssumeHooksFollowRulesOfReact = + config?.enableAssumeHooksFollowRulesOfReact ?? false; this.#contextIdentifiers = contextIdentifiers; } @@ -98,12 +139,7 @@ export class Environment { if (isHookName(name)) { return { kind: "Hook", - definition: { - kind: "Custom", - name, - effectKind: Effect.Mutate, - valueKind: ValueKind.Mutable, - }, + definition: null, }; } else { log(() => `Undefined global '${name}'`); diff --git a/compiler/forget/src/HIR/Types.ts b/compiler/forget/src/HIR/Types.ts index 0f272ceab7..14083d107f 100644 --- a/compiler/forget/src/HIR/Types.ts +++ b/compiler/forget/src/HIR/Types.ts @@ -20,7 +20,7 @@ export type Type = export type PrimitiveType = { kind: "Primitive" }; export type HookType = { kind: "Hook"; - definition: Hook; + definition: Hook | null; }; /** diff --git a/compiler/forget/src/Inference/DropMemoCalls.ts b/compiler/forget/src/Inference/DropMemoCalls.ts index b9c0619056..4b1669f7d3 100644 --- a/compiler/forget/src/Inference/DropMemoCalls.ts +++ b/compiler/forget/src/Inference/DropMemoCalls.ts @@ -14,7 +14,7 @@ export default function (func: HIRFunction): void { case "CallExpression": { if (isHookType(instr.value.callee.identifier)) { const name = (instr.value.callee.identifier.type as HookType) - .definition.name; + .definition?.name; if (name === "useMemo") { const [fn] = instr.value.args; diff --git a/compiler/forget/src/Inference/InferReferenceEffects.ts b/compiler/forget/src/Inference/InferReferenceEffects.ts index b368a135bb..c58b8fe872 100644 --- a/compiler/forget/src/Inference/InferReferenceEffects.ts +++ b/compiler/forget/src/Inference/InferReferenceEffects.ts @@ -708,13 +708,18 @@ function inferBlock( continue; } case "CallExpression": { - const hook = - instrValue.callee.identifier.type.kind === "Hook" - ? instrValue.callee.identifier.type.definition - : null; - if (hook !== null) { - effectKind = hook.effectKind; - valueKind = hook.valueKind; + if (instrValue.callee.identifier.type.kind === "Hook") { + const definition = instrValue.callee.identifier.type.definition; + if (definition !== null) { + effectKind = definition.effectKind; + valueKind = definition.valueKind; + } else if (env.enableAssumeHooksFollowRulesOfReact) { + effectKind = Effect.Freeze; + valueKind = ValueKind.Frozen; + } else { + effectKind = Effect.Mutate; + valueKind = ValueKind.Mutable; + } break; } diff --git a/compiler/forget/src/__tests__/compiler-test.ts b/compiler/forget/src/__tests__/compiler-test.ts index 27ba569653..a904127cb5 100644 --- a/compiler/forget/src/__tests__/compiler-test.ts +++ b/compiler/forget/src/__tests__/compiler-test.ts @@ -60,6 +60,9 @@ describe("React Forget", () => { validateHooksUsage: true, inlineUseMemo: options.environment?.inlineUseMemo ?? false, enableFunctionCallSignatureOptimizations: true, + enableAssumeHooksFollowRulesOfReact: + options.environment?.enableAssumeHooksFollowRulesOfReact ?? + false, }, logger: null, gating: options.gating, diff --git a/compiler/forget/src/__tests__/fixtures/compiler/immutable-hooks.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/immutable-hooks.expect.md new file mode 100644 index 0000000000..a4f065bb82 --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/immutable-hooks.expect.md @@ -0,0 +1,48 @@ + +## Input + +```javascript +// @enableAssumeHooksFollowRulesOfReact true +function Component(props) { + const x = {}; + // In enableAssumeHooksFollowRulesOfReact mode hooks freeze their inputs and return frozen values + const y = useFoo(x); + // Thus both x and y are frozen here, and x can be independently memoized + bar(x, y); + return [x, y]; +} + +``` + +## Code + +```javascript +import { unstable_useMemoCache as useMemoCache } from "react"; // @enableAssumeHooksFollowRulesOfReact true +function Component(props) { + const $ = useMemoCache(3); + let t0; + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + t0 = {}; + $[0] = t0; + } else { + t0 = $[0]; + } + const x = t0; + + const y = useFoo(x); + + bar(x, y); + const c_1 = $[1] !== y; + let t1; + if (c_1) { + t1 = [x, y]; + $[1] = y; + $[2] = t1; + } else { + t1 = $[2]; + } + return t1; +} + +``` + \ No newline at end of file diff --git a/compiler/forget/src/__tests__/fixtures/compiler/immutable-hooks.js b/compiler/forget/src/__tests__/fixtures/compiler/immutable-hooks.js new file mode 100644 index 0000000000..b4103fa5da --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/immutable-hooks.js @@ -0,0 +1,9 @@ +// @enableAssumeHooksFollowRulesOfReact true +function Component(props) { + const x = {}; + // In enableAssumeHooksFollowRulesOfReact mode hooks freeze their inputs and return frozen values + const y = useFoo(x); + // Thus both x and y are frozen here, and x can be independently memoized + bar(x, y); + return [x, y]; +} diff --git a/compiler/forget/src/__tests__/test-utils/generateTestsFromFixtures.ts b/compiler/forget/src/__tests__/test-utils/generateTestsFromFixtures.ts index 6256789e58..d7c05a95d5 100644 --- a/compiler/forget/src/__tests__/test-utils/generateTestsFromFixtures.ts +++ b/compiler/forget/src/__tests__/test-utils/generateTestsFromFixtures.ts @@ -102,6 +102,7 @@ export default function generateTestsFromFixtures( let inlineUseMemo = true; let panicOnBailout = true; let memoizeJsxElements = true; + let enableAssumeHooksFollowRulesOfReact = false; if (inputFile != null) { input = fs.readFileSync(inputFile, "utf8"); @@ -145,13 +146,23 @@ export default function generateTestsFromFixtures( if (lines[0]!.indexOf("@memoizeJsxElements false") !== -1) { memoizeJsxElements = false; } + if ( + lines[0]!.indexOf("@enableAssumeHooksFollowRulesOfReact true") !== + -1 + ) { + enableAssumeHooksFollowRulesOfReact = true; + } } testCommand(basename, () => { let receivedOutput; if (input !== null) { receivedOutput = transform(input, basename, { - environment: { inlineUseMemo, memoizeJsxElements }, + environment: { + inlineUseMemo, + memoizeJsxElements, + enableAssumeHooksFollowRulesOfReact, + }, logger: null, debug, enableOnlyOnUseForgetDirective,