@samitouri / QOS-React / commits / 4ce22ebcfe

Option to assume hooks follow the rules

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.

Joe Savona committed May 19, 2023 at 11:05 UTC 4ce22ebcfeb264b6b20e5a5b6f7fc16585a0f907
9 files changed +135 -18
compiler/forget/packages/snap/src/compiler-worker.ts
+5
@@ -93,6 +93,7 @@ export async function compile(
93 let instrumentForget = null;
94 let panicOnBailout = true;
95 let memoizeJsxElements = true;
96 + let enableAssumeHooksFollowRulesOfReact = false;
97 if (firstLine.indexOf("@forgetDirective") !== -1) {
98 enableOnlyOnUseForgetDirective = true;
99 }
@@ -120,6 +121,9 @@ export async function compile(
121 if (firstLine.indexOf("@memoizeJsxElements false") !== -1) {
122 memoizeJsxElements = false;
123 }
124 + if (firstLine.indexOf("@enableAssumeHooksFollowRulesOfReact true") !== -1) {
125 + enableAssumeHooksFollowRulesOfReact = true;
126 + }
127
128 const language = parseLanguage(firstLine);
129
@@ -139,6 +143,7 @@ export async function compile(
143 ]),
144 validateHooksUsage: true,
145 enableFunctionCallSignatureOptimizations: true,
146 + enableAssumeHooksFollowRulesOfReact,
147 inlineUseMemo: true,
148 memoizeJsxElements,
149 },
compiler/forget/src/HIR/Environment.ts
+44 -8
@@ -17,12 +17,10 @@ import {
17 import {
18 BlockId,
19 BuiltInType,
20 - Effect,
20 FunctionType,
21 IdentifierId,
22 ObjectType,
23 PolyType,
25 - ValueKind,
24 makeBlockId,
25 makeIdentifierId,
26 } from "./HIR";
@@ -38,10 +36,50 @@ import { FunctionSignature, ShapeRegistry } from "./ObjectShape";
36 // missing some recursive Object / Function shapeIds
37 export type EnvironmentConfig = Partial<{
38 customHooks: Map<string, Hook>;
39 +
40 + /**
41 + * Enable memoization of JSX elements in addition to other types of values. When disabled,
42 + * other types (objects, arrays, call expressions, etc) are memoized, but not known JSX
43 + * values.
44 + *
45 + * Defaults to true
46 + */
47 memoizeJsxElements: boolean;
48 +
49 + /**
50 + * Enable validation of hooks to partially check that the component honors the rules of hooks.
51 + * When disabled, the component is assumed to follow the rules (though the Babel plugin looks
52 + * for suppressions of the lint rule).
53 + *
54 + * Defaults to false
55 + */
56 validateHooksUsage: boolean;
57 +
58 + /**
59 + * Enable inlining of `useMemo()` function expressions so that they can be more optimally
60 + * compiled.
61 + *
62 + * Defaults to false
63 + */
64 inlineUseMemo: boolean;
65 +
66 + /**
67 + * Enable optimizations based on the signature of (non-method) built-in function calls.
68 + *
69 + * Defaults to false
70 + */
71 enableFunctionCallSignatureOptimizations: boolean;
72 +
73 + /**
74 + * When enabled, the compiler assumes that hooks follow the Rules of React:
75 + * - Hooks may memoize computation based on any of their parameters, thus
76 + * any arguments to a hook are assumed frozen after calling the hook.
77 + * - Hooks may memoize the result they return, thus the return value is
78 + * assumed frozen.
79 +
80 + * Defaults to false
81 + */
82 + enableAssumeHooksFollowRulesOfReact: boolean;
83 }>;
84
85 export class Environment {
@@ -51,6 +89,7 @@ export class Environment {
89 #nextBlock: number = 0;
90 validateHooksUsage: boolean;
91 enableFunctionCallSignatureOptimizations: boolean;
92 + enableAssumeHooksFollowRulesOfReact: boolean;
93 #contextIdentifiers: Set<t.Identifier>;
94
95 constructor(
@@ -77,6 +116,8 @@ export class Environment {
116 this.validateHooksUsage = config?.validateHooksUsage ?? false;
117 this.enableFunctionCallSignatureOptimizations =
118 config?.enableFunctionCallSignatureOptimizations ?? false;
119 + this.enableAssumeHooksFollowRulesOfReact =
120 + config?.enableAssumeHooksFollowRulesOfReact ?? false;
121 this.#contextIdentifiers = contextIdentifiers;
122 }
123
@@ -98,12 +139,7 @@ export class Environment {
139 if (isHookName(name)) {
140 return {
141 kind: "Hook",
101 - definition: {
102 - kind: "Custom",
103 - name,
104 - effectKind: Effect.Mutate,
105 - valueKind: ValueKind.Mutable,
106 - },
142 + definition: null,
143 };
144 } else {
145 log(() => `Undefined global '${name}'`);
compiler/forget/src/HIR/Types.ts
+1 -1
@@ -20,7 +20,7 @@ export type Type =
20 export type PrimitiveType = { kind: "Primitive" };
21 export type HookType = {
22 kind: "Hook";
23 - definition: Hook;
23 + definition: Hook | null;
24 };
25
26 /**
compiler/forget/src/Inference/DropMemoCalls.ts
+1 -1
@@ -14,7 +14,7 @@ export default function (func: HIRFunction): void {
14 case "CallExpression": {
15 if (isHookType(instr.value.callee.identifier)) {
16 const name = (instr.value.callee.identifier.type as HookType)
17 - .definition.name;
17 + .definition?.name;
18 if (name === "useMemo") {
19 const [fn] = instr.value.args;
20
compiler/forget/src/Inference/InferReferenceEffects.ts
+12 -7
@@ -708,13 +708,18 @@ function inferBlock(
708 continue;
709 }
710 case "CallExpression": {
711 - const hook =
712 - instrValue.callee.identifier.type.kind === "Hook"
713 - ? instrValue.callee.identifier.type.definition
714 - : null;
715 - if (hook !== null) {
716 - effectKind = hook.effectKind;
717 - valueKind = hook.valueKind;
711 + if (instrValue.callee.identifier.type.kind === "Hook") {
712 + const definition = instrValue.callee.identifier.type.definition;
713 + if (definition !== null) {
714 + effectKind = definition.effectKind;
715 + valueKind = definition.valueKind;
716 + } else if (env.enableAssumeHooksFollowRulesOfReact) {
717 + effectKind = Effect.Freeze;
718 + valueKind = ValueKind.Frozen;
719 + } else {
720 + effectKind = Effect.Mutate;
721 + valueKind = ValueKind.Mutable;
722 + }
723 break;
724 }
725
compiler/forget/src/__tests__/compiler-test.ts
+3
@@ -60,6 +60,9 @@ describe("React Forget", () => {
60 validateHooksUsage: true,
61 inlineUseMemo: options.environment?.inlineUseMemo ?? false,
62 enableFunctionCallSignatureOptimizations: true,
63 + enableAssumeHooksFollowRulesOfReact:
64 + options.environment?.enableAssumeHooksFollowRulesOfReact ??
65 + false,
66 },
67 logger: null,
68 gating: options.gating,
compiler/forget/src/__tests__/fixtures/compiler/immutable-hooks.expect.md new
+48
@@ -0,0 +1,48 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @enableAssumeHooksFollowRulesOfReact true
6 +function Component(props) {
7 + const x = {};
8 + // In enableAssumeHooksFollowRulesOfReact mode hooks freeze their inputs and return frozen values
9 + const y = useFoo(x);
10 + // Thus both x and y are frozen here, and x can be independently memoized
11 + bar(x, y);
12 + return [x, y];
13 +}
14 +
15 +```
16 +
17 +## Code
18 +
19 +```javascript
20 +import { unstable_useMemoCache as useMemoCache } from "react"; // @enableAssumeHooksFollowRulesOfReact true
21 +function Component(props) {
22 + const $ = useMemoCache(3);
23 + let t0;
24 + if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
25 + t0 = {};
26 + $[0] = t0;
27 + } else {
28 + t0 = $[0];
29 + }
30 + const x = t0;
31 +
32 + const y = useFoo(x);
33 +
34 + bar(x, y);
35 + const c_1 = $[1] !== y;
36 + let t1;
37 + if (c_1) {
38 + t1 = [x, y];
39 + $[1] = y;
40 + $[2] = t1;
41 + } else {
42 + t1 = $[2];
43 + }
44 + return t1;
45 +}
46 +
47 +```
48 +
\ No newline at end of file
compiler/forget/src/__tests__/fixtures/compiler/immutable-hooks.js new
+9
@@ -0,0 +1,9 @@
1 +// @enableAssumeHooksFollowRulesOfReact true
2 +function Component(props) {
3 + const x = {};
4 + // In enableAssumeHooksFollowRulesOfReact mode hooks freeze their inputs and return frozen values
5 + const y = useFoo(x);
6 + // Thus both x and y are frozen here, and x can be independently memoized
7 + bar(x, y);
8 + return [x, y];
9 +}
compiler/forget/src/__tests__/test-utils/generateTestsFromFixtures.ts
+12 -1
@@ -102,6 +102,7 @@ export default function generateTestsFromFixtures(
102 let inlineUseMemo = true;
103 let panicOnBailout = true;
104 let memoizeJsxElements = true;
105 + let enableAssumeHooksFollowRulesOfReact = false;
106
107 if (inputFile != null) {
108 input = fs.readFileSync(inputFile, "utf8");
@@ -145,13 +146,23 @@ export default function generateTestsFromFixtures(
146 if (lines[0]!.indexOf("@memoizeJsxElements false") !== -1) {
147 memoizeJsxElements = false;
148 }
149 + if (
150 + lines[0]!.indexOf("@enableAssumeHooksFollowRulesOfReact true") !==
151 + -1
152 + ) {
153 + enableAssumeHooksFollowRulesOfReact = true;
154 + }
155 }
156
157 testCommand(basename, () => {
158 let receivedOutput;
159 if (input !== null) {
160 receivedOutput = transform(input, basename, {
154 - environment: { inlineUseMemo, memoizeJsxElements },
161 + environment: {
162 + inlineUseMemo,
163 + memoizeJsxElements,
164 + enableAssumeHooksFollowRulesOfReact,
165 + },
166 logger: null,
167 debug,
168 enableOnlyOnUseForgetDirective,