@samitouri / QOS-React / commits / 81695f62c2

[refactor] Refactor Memoize to two instructions: Start and Finish

--- Previously, we always emitted `Memoize dep` instructions after the function expression literal and depslist instructions ```js // source useManualMemo(() => {...}, [arg]) // lowered $0 = FunctionExpression(...) $1 = LoadLocal (arg) $2 = ArrayExpression [$1] $3 = Memoize (arg) $4 = Call / LoadLocal $5 = Memoize $4 ``` Now, we insert `Memoize dep` before the corresponding function expression literal: ```js // lowered $0 = StartMemoize (arg) <---- this moved up! $1 = FunctionExpression(...) $2 = LoadLocal (arg) $3 = ArrayExpression [$2] $4 = Call / LoadLocal $5 = FinishMemoize $4 ``` Design considerations: - #2663 needs to understand which lowered instructions belong to a manual memoization block, so we need to emit `StartMemoize` instructions before the `useMemo/useCallback` function argument, which contains relevant memoized instructions - we choose to insert StartMemoize instructions to (1) avoid unsafe instruction reordering of source and (2) to ensure that Forget output does not change when enabling validation This PR only renames `Memoize` -> `Start/FinishMemoize` and hoists `StartMemoize` as described. The latter may help with stricter validation for `useCallback`s, although testing is left to the next PR. #2663 contains all validation changes

Mofei Zhang committed Mar 18, 2024 at 12:09 UTC 81695f62c2b544b2f5c91cc032a01ab1a52fcef9
16 files changed +453 -324
compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts
+35 -12
@@ -528,6 +528,13 @@ export type Instruction = {
528 loc: SourceLocation;
529 };
530
531 +export type TInstruction<T extends InstructionValue> = {
532 + id: InstructionId;
533 + lvalue: Place;
534 + value: T;
535 + loc: SourceLocation;
536 +};
537 +
538 export type LValue = {
539 place: Place;
540 kind: InstructionKind;
@@ -625,6 +632,17 @@ export type Phi = {
632 type: Type;
633 };
634
635 +export type StartMemoize = {
636 + kind: "StartMemoize";
637 + deps: Array<Place>;
638 + loc: SourceLocation;
639 +};
640 +export type FinishMemoize = {
641 + kind: "FinishMemoize";
642 + decl: Place;
643 + loc: SourceLocation;
644 +};
645 +
646 /*
647 * Forget currently does not handle MethodCall correctly in
648 * all cases. Specifically, we do not bind the receiver and method property
@@ -657,6 +675,12 @@ export type CallExpression = {
675 typeArguments?: Array<t.FlowType>;
676 };
677
678 +export type LoadLocal = {
679 + kind: "LoadLocal";
680 + place: Place;
681 + loc: SourceLocation;
682 +};
683 +
684 /*
685 * The value of a given instruction. Note that values are not recursive: complex
686 * values such as objects or arrays are always defined by instructions to define
@@ -667,11 +691,7 @@ export type CallExpression = {
691 */
692
693 export type InstructionValue =
670 - | {
671 - kind: "LoadLocal";
672 - place: Place;
673 - loc: SourceLocation;
674 - }
694 + | LoadLocal
695 | {
696 kind: "LoadContext";
697 place: Place;
@@ -773,12 +793,7 @@ export type InstructionValue =
793 loc: SourceLocation;
794 }
795 // load `object.property`
776 - | {
777 - kind: "PropertyLoad";
778 - object: Place;
779 - property: string;
780 - loc: SourceLocation;
781 - }
796 + | PropertyLoad
797 // `delete object.property`
798 | {
799 kind: "PropertyDelete";
@@ -873,7 +888,8 @@ export type InstructionValue =
888 * during codegen. It can't be pruned during DCE because we need to preserve the
889 * instruction so it can be visible in InferReferenceEffects.
890 */
876 - | { kind: "Memoize"; value: Place; loc: SourceLocation }
891 + | StartMemoize
892 + | FinishMemoize
893 /*
894 * Catch-all for statements such as type imports, nested class declarations, etc
895 * which are not directly represented, but included for completeness and to allow
@@ -929,6 +945,13 @@ export type Primitive = {
945
946 export type JSXText = { kind: "JSXText"; value: string; loc: SourceLocation };
947
948 +export type PropertyLoad = {
949 + kind: "PropertyLoad";
950 + object: Place;
951 + property: string;
952 + loc: SourceLocation;
953 +};
954 +
955 export type LoadGlobal = {
956 kind: "LoadGlobal";
957 name: string;
compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts
+8 -2
@@ -600,8 +600,14 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
600 } ${printPlace(instrValue.value)}`;
601 break;
602 }
603 - case "Memoize": {
604 - value = `Memoize ${printPlace(instrValue.value)}`;
603 + case "StartMemoize": {
604 + value = `StartMemoize deps=${instrValue.deps.map((dep) =>
605 + printPlace(dep)
606 + )}`;
607 + break;
608 + }
609 + case "FinishMemoize": {
610 + value = `FinishMemoize decl=${printPlace(instrValue.decl)}`;
611 break;
612 }
613 case "ReactiveFunctionValue": {
compiler/packages/babel-plugin-react-forget/src/HIR/visitors.ts
+16 -4
@@ -216,8 +216,14 @@ export function* eachInstructionValueOperand(
216 yield instrValue.value;
217 break;
218 }
219 - case "Memoize": {
220 - yield instrValue.value;
219 + case "StartMemoize": {
220 + for (const dep of instrValue.deps) {
221 + yield dep;
222 + }
223 + break;
224 + }
225 + case "FinishMemoize": {
226 + yield instrValue.decl;
227 break;
228 }
229 case "Debugger":
@@ -521,8 +527,14 @@ export function mapInstructionValueOperands(
527 instrValue.value = fn(instrValue.value);
528 break;
529 }
524 - case "Memoize": {
525 - instrValue.value = fn(instrValue.value);
530 + case "StartMemoize": {
531 + for (let i = 0; i < instrValue.deps.length; i++) {
532 + instrValue.deps[i] = fn(instrValue.deps[i]);
533 + }
534 + break;
535 + }
536 + case "FinishMemoize": {
537 + instrValue.decl = fn(instrValue.decl);
538 break;
539 }
540 case "Debugger":
compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts
+305 -256
@@ -5,22 +5,210 @@
5 * LICENSE file in the root directory of this source tree.
6 */
7
8 -import { CompilerError } from "..";
8 +import { CompilerError, SourceLocation } from "..";
9 import {
10 + CallExpression,
11 Effect,
12 + Environment,
13 + FinishMemoize,
14 FunctionExpression,
15 HIRFunction,
16 IdentifierId,
17 Instruction,
18 + InstructionId,
19 + LoadGlobal,
20 + LoadLocal,
21 + MethodCall,
22 Place,
23 + PropertyLoad,
24 SpreadPattern,
25 + StartMemoize,
26 + TInstruction,
27 getHookKindForType,
28 makeInstructionId,
29 } from "../HIR";
30 import { createTemporaryPlace, markInstructionIds } from "../HIR/HIRBuilder";
21 -import { HookKind } from "../HIR/ObjectShape";
31 import { eachInstructionValueOperand } from "../HIR/visitors";
32
33 +type ManualMemoCallee = {
34 + kind: "useMemo" | "useCallback";
35 + loadInstr: TInstruction<LoadGlobal> | TInstruction<PropertyLoad>;
36 +};
37 +
38 +type IdentifierSidemap = {
39 + functions: Map<IdentifierId, TInstruction<FunctionExpression>>;
40 + manualMemos: Map<IdentifierId, ManualMemoCallee>;
41 + react: Set<IdentifierId>;
42 +};
43 +
44 +function collectTemporaries(
45 + instr: Instruction,
46 + env: Environment,
47 + sidemap: IdentifierSidemap
48 +): void {
49 + const { value } = instr;
50 + switch (value.kind) {
51 + case "FunctionExpression": {
52 + sidemap.functions.set(
53 + instr.lvalue.identifier.id,
54 + instr as TInstruction<FunctionExpression>
55 + );
56 + break;
57 + }
58 + case "LoadGlobal": {
59 + const global = env.getGlobalDeclaration(value.name);
60 + const hookKind = global !== null ? getHookKindForType(env, global) : null;
61 + const lvalId = instr.lvalue.identifier.id;
62 + if (hookKind === "useMemo" || hookKind === "useCallback") {
63 + sidemap.manualMemos.set(lvalId, {
64 + kind: hookKind,
65 + loadInstr: instr as TInstruction<LoadGlobal>,
66 + });
67 + } else if (value.name === "React") {
68 + sidemap.react.add(lvalId);
69 + }
70 + break;
71 + }
72 + case "PropertyLoad": {
73 + if (sidemap.react.has(value.object.identifier.id)) {
74 + if (value.property === "useMemo" || value.property === "useCallback") {
75 + sidemap.manualMemos.set(instr.lvalue.identifier.id, {
76 + kind: value.property,
77 + loadInstr: instr as TInstruction<PropertyLoad>,
78 + });
79 + }
80 + }
81 + break;
82 + }
83 + }
84 +}
85 +
86 +function makeManualMemoizationMarkers(
87 + fnExpr: Place,
88 + env: Environment,
89 + depsList: Array<Place>,
90 + memoDecl: Place
91 +): [TInstruction<StartMemoize>, TInstruction<FinishMemoize>] {
92 + return [
93 + {
94 + id: makeInstructionId(0),
95 + lvalue: createTemporaryPlace(env),
96 + value: {
97 + kind: "StartMemoize",
98 + /*
99 + * Use deps list from source instead of inferred deps
100 + * as dependencies
101 + */
102 + deps: depsList,
103 + loc: fnExpr.loc,
104 + },
105 + loc: fnExpr.loc,
106 + },
107 + {
108 + id: makeInstructionId(0),
109 + lvalue: createTemporaryPlace(env),
110 + value: {
111 + kind: "FinishMemoize",
112 + decl: { ...memoDecl },
113 + loc: fnExpr.loc,
114 + },
115 + loc: fnExpr.loc,
116 + },
117 + ];
118 +}
119 +
120 +function getManualMemoizationReplacement(
121 + fn: Place,
122 + loc: SourceLocation,
123 + kind: "useMemo" | "useCallback"
124 +): LoadLocal | CallExpression {
125 + if (kind === "useMemo") {
126 + /*
127 + * Replace the hook callee with the fn arg.
128 + *
129 + * before:
130 + * $1 = LoadGlobal useMemo // load the useMemo global
131 + * $2 = FunctionExpression ... // memo function
132 + * $3 = ArrayExpression [ ... ] // deps array
133 + * $4 = Call $1 ($2, $3 ) // invoke useMemo w fn and deps
134 + *
135 + * after:
136 + * $1 = LoadGlobal useMemo // load the useMemo global (dead code)
137 + * $2 = FunctionExpression ... // memo function
138 + * $3 = ArrayExpression [ ... ] // deps array (dead code)
139 + * $4 = Call $2 () // invoke the memo function itself
140 + *
141 + * Note that a later pass (InlineImmediatelyInvokedFunctionExpressions) will
142 + * inline the useMemo callback along with any other immediately invoked IIFEs.
143 + */
144 + return {
145 + kind: "CallExpression",
146 + callee: fn,
147 + /*
148 + * Drop the args, including the deps array which DCE will remove
149 + * later.
150 + */
151 + args: [],
152 + loc,
153 + };
154 + } else {
155 + /*
156 + * Instead of a Call, just alias the callback directly.
157 + *
158 + * before:
159 + * $1 = LoadGlobal useCallback
160 + * $2 = FunctionExpression ... // the callback being memoized
161 + * $3 = ArrayExpression ... // deps array
162 + * $4 = Call $1 ( $2, $3 ) // invoke useCallback
163 + *
164 + * after:
165 + * $1 = LoadGlobal useCallback // dead code
166 + * $2 = FunctionExpression ... // the callback being memoized
167 + * $3 = ArrayExpression ... // deps array (dead code)
168 + * $4 = LoadLocal $2 // reference the function
169 + */
170 + return {
171 + kind: "LoadLocal",
172 + place: {
173 + kind: "Identifier",
174 + identifier: fn.identifier,
175 + effect: Effect.Unknown,
176 + reactive: false,
177 + loc,
178 + },
179 + loc,
180 + };
181 + }
182 +}
183 +
184 +function extractManualMemoizationArgs(
185 + instr: TInstruction<CallExpression> | TInstruction<MethodCall>,
186 + kind: "useCallback" | "useMemo"
187 +): {
188 + fnPlace: Place;
189 +} {
190 + const [fnPlace] = instr.value.args as Array<
191 + Place | SpreadPattern | undefined
192 + >;
193 + if (fnPlace == null) {
194 + CompilerError.throwInvalidReact({
195 + reason: `Expected ${kind} call to pass a callback function`,
196 + loc: instr.value.loc,
197 + suggestions: null,
198 + });
199 + }
200 + if (fnPlace?.kind !== "Identifier") {
201 + CompilerError.throwInvalidReact({
202 + reason: `Unexpected arguments to ${kind} call`,
203 + loc: instr.value.loc,
204 + suggestions: null,
205 + });
206 + }
207 + return {
208 + fnPlace,
209 + };
210 +}
211 +
212 /*
213 * Removes manual memoization using the `useMemo` and `useCallback` APIs. This pass is designed
214 * to compose with InlineImmediatelyInvokedFunctionExpressions, and needs to run prior to entering
@@ -30,274 +218,135 @@ import { eachInstructionValueOperand } from "../HIR/visitors";
218 * eg `React.useMemo()`.
219 */
220 export function dropManualMemoization(func: HIRFunction): void {
33 - const functions = new Map<IdentifierId, FunctionExpression>();
34 - const hooks = new Map<IdentifierId, HookKind>();
35 - const react = new Set<IdentifierId>();
36 - let hasChanges = false;
221 + const isValidationEnabled =
222 + func.env.config.validatePreserveExistingMemoizationGuarantees ||
223 + func.env.config.enablePreserveExistingMemoizationGuarantees;
224 + const sidemap: IdentifierSidemap = {
225 + functions: new Map(),
226 + manualMemos: new Map(),
227 + react: new Set(),
228 + };
229 +
230 + /**
231 + * Phase 1:
232 + * - Overwrite manual memoization from
233 + * CallExpression callee="useMemo/Callback", args=[fnArg, depslist])
234 + * to either
235 + * CallExpression callee=fnArg
236 + * LoadLocal fnArg
237 + * - (if validation is enabled) collect manual memoization markers
238 + */
239 + const queuedInserts: Map<
240 + InstructionId,
241 + {
242 + kind: "before" | "after";
243 + value: TInstruction<StartMemoize> | TInstruction<FinishMemoize>;
244 + }
245 + > = new Map();
246 for (const [_, block] of func.body.blocks) {
38 - let nextInstructions: Array<Instruction> | null = null;
247 for (let i = 0; i < block.instructions.length; i++) {
248 const instr = block.instructions[i]!;
41 - switch (instr.value.kind) {
42 - case "FunctionExpression": {
43 - functions.set(instr.lvalue.identifier.id, instr.value);
44 - break;
45 - }
46 - case "LoadGlobal": {
47 - const global = func.env.getGlobalDeclaration(instr.value.name);
48 - const hookKind =
49 - global !== null ? getHookKindForType(func.env, global) : null;
50 - if (hookKind === "useMemo" || hookKind === "useCallback") {
51 - hooks.set(instr.lvalue.identifier.id, hookKind);
52 - } else if (instr.value.name === "React") {
53 - react.add(instr.lvalue.identifier.id);
54 - }
55 - break;
56 - }
57 - case "PropertyLoad": {
58 - if (react.has(instr.value.object.identifier.id)) {
59 - if (
60 - instr.value.property === "useMemo" ||
61 - instr.value.property === "useCallback"
62 - ) {
63 - hooks.set(instr.lvalue.identifier.id, instr.value.property);
64 - }
65 - }
66 - break;
67 - }
68 - case "MethodCall":
69 - case "CallExpression": {
70 - const id =
71 - instr.value.kind === "CallExpression"
72 - ? instr.value.callee.identifier.id
73 - : instr.value.property.identifier.id;
74 - const hookKind = hooks.get(id);
75 - if (hookKind != null) {
76 - if (hookKind === "useMemo") {
77 - const [fn] = instr.value.args as Array<
78 - Place | SpreadPattern | undefined
79 - >;
80 - if (fn == null) {
81 - CompilerError.throwInvalidReact({
82 - reason: "Expected useMemo call to pass a callback function",
83 - loc: instr.loc,
84 - suggestions: null,
85 - });
86 - }
87 - /*
88 - * Replace the hook callee with the fn arg.
89 - *
90 - * before:
91 - * $1 = LoadGlobal useMemo // load the useMemo global
92 - * $2 = FunctionExpression ... // memo function
93 - * $3 = ArrayExpression [ ... ] // deps array
94 - * $4 = Call $1 ($2, $3 ) // invoke useMemo w fn and deps
95 - *
96 - * after:
97 - * $1 = LoadGlobal useMemo // load the useMemo global (dead code)
98 - * $2 = FunctionExpression ... // memo function
99 - * $3 = ArrayExpression [ ... ] // deps array (dead code)
100 - * $4 = Call $2 () // invoke the memo function itself
101 - *
102 - * Note that a later pass (InlineImmediatelyInvokedFunctionExpressions) will
103 - * inline the useMemo callback along with any other immediately invoked IIFEs.
104 - */
105 - if (fn.kind === "Identifier") {
106 - instr.value = {
107 - kind: "CallExpression",
108 - callee: fn,
109 - /*
110 - * Drop the args, including the deps array which DCE will remove
111 - * later.
112 - */
113 - args: [],
114 - loc: instr.value.loc,
115 - };
249 + if (
250 + instr.value.kind === "CallExpression" ||
251 + instr.value.kind === "MethodCall"
252 + ) {
253 + const id =
254 + instr.value.kind === "CallExpression"
255 + ? instr.value.callee.identifier.id
256 + : instr.value.property.identifier.id;
257
117 - if (
118 - func.env.config.enablePreserveExistingMemoizationGuarantees ||
119 - func.env.config.validatePreserveExistingMemoizationGuarantees
120 - ) {
121 - /**
122 - * When this flag is enabled we also compile in a 'Memoize' instruction
123 - * to preserve the intended memoization boundary:
124 - *
125 - * Normal output:
126 - * $1 = LoadGlobal useMemo // load the useMemo global (dead code)
127 - * $2 = FunctionExpression ... // memo function
128 - * $3 = ArrayExpression [ ... ] // deps array (dead code)
129 - * $4 = Call $2 () // invoke the memo function itself
130 - *
131 - * Output w flag enabled:
132 - * $1 = LoadGlobal useMemo // load the useMemo global (dead code)
133 - * $2 = FunctionExpression ... // memo function
134 - * $3 = ArrayExpression [ ... ] // deps array (dead code)
135 - * .. = Memoize ... // memoize dependencies
136 - * $4 = Call $2 () // invoke the memo function itself
137 - * .. = Memoize $4 // preserve memo information
138 - *
139 - * Note that Memoize does not produce a result and is called for its side
140 - * effects only.
141 - */
142 - nextInstructions =
143 - nextInstructions ?? block.instructions.slice(0, i);
144 -
145 - const functionExpression = functions.get(fn.identifier.id);
146 - if (functionExpression !== undefined) {
147 - for (const operand of eachInstructionValueOperand(
148 - functionExpression
149 - )) {
150 - const temp = createTemporaryPlace(func.env);
151 - nextInstructions.push({
152 - id: makeInstructionId(0),
153 - lvalue: temp,
154 - value: {
155 - kind: "Memoize",
156 - value: { ...operand },
157 - loc: instr.loc,
158 - },
159 - loc: instr.loc,
160 - });
161 - }
162 - }
163 -
164 - nextInstructions.push(instr);
165 -
166 - const temp = createTemporaryPlace(func.env);
167 - nextInstructions.push({
168 - id: makeInstructionId(0),
169 - lvalue: temp,
170 - value: {
171 - kind: "Memoize",
172 - value: { ...instr.lvalue },
173 - loc: instr.loc,
174 - },
175 - loc: instr.loc,
176 - });
177 - continue;
178 - }
179 - }
180 - } else if (hookKind === "useCallback") {
181 - const [fn] = instr.value.args as Array<
182 - Place | SpreadPattern | undefined
183 - >;
184 - if (fn == null) {
185 - CompilerError.throwInvalidReact({
186 - reason: "Expected useMemo call to pass a callback function",
187 - loc: instr.loc,
188 - suggestions: null,
189 - });
190 - }
191 -
192 - /*
193 - * Instead of a Call, just alias the callback directly.
194 - *
195 - * before:
196 - * $1 = LoadGlobal useCallback
197 - * $2 = FunctionExpression ... // the callback being memoized
198 - * $3 = ArrayExpression ... // deps array
199 - * $4 = Call $1 ( $2, $3 ) // invoke useCallback
200 - *
201 - * after:
202 - * $1 = LoadGlobal useCallback // dead code
203 - * $2 = FunctionExpression ... // the callback being memoized
204 - * $3 = ArrayExpression ... // deps array (dead code)
205 - * $4 = LoadLocal $2 // reference the function
206 - */
207 - if (fn.kind === "Identifier") {
208 - instr.value = {
209 - kind: "LoadLocal",
210 - place: {
258 + const manualMemo = sidemap.manualMemos.get(id);
259 + if (manualMemo != null) {
260 + const { fnPlace } = extractManualMemoizationArgs(
261 + instr as TInstruction<CallExpression> | TInstruction<MethodCall>,
262 + manualMemo.kind
263 + );
264 + instr.value = getManualMemoizationReplacement(
265 + fnPlace,
266 + instr.value.loc,
267 + manualMemo.kind
268 + );
269 + if (isValidationEnabled) {
270 + const inlineMemoFn = sidemap.functions.get(fnPlace.identifier.id);
271 + if (inlineMemoFn == null) {
272 + CompilerError.throwInvalidReact({
273 + reason:
274 + "DepsValidation: Expected function literal as manual memoization callback",
275 + suggestions: [],
276 + loc: fnPlace.loc,
277 + });
278 + }
279 + const memoDecl: Place =
280 + manualMemo.kind === "useMemo"
281 + ? instr.lvalue
282 + : {
283 kind: "Identifier",
212 - identifier: fn.identifier,
284 + identifier: fnPlace.identifier,
285 effect: Effect.Unknown,
286 reactive: false,
215 - loc: instr.value.loc,
216 - },
217 - loc: instr.value.loc,
218 - };
219 - if (
220 - func.env.config.enablePreserveExistingMemoizationGuarantees ||
221 - func.env.config.validatePreserveExistingMemoizationGuarantees
222 - ) {
223 - nextInstructions =
224 - nextInstructions ?? block.instructions.slice(0, i);
225 - /**
226 - * With the flag enabled the output changes to use a Memoize instruction instead
227 - * a loadlocal to load the function expression into the original temporary:
228 - *
229 - * Normal output:
230 - * $1 = LoadGlobal useCallback // dead code
231 - * $2 = FunctionExpression ... // the callback being memoized
232 - * $3 = ArrayExpression ... // deps array (dead code)
233 - * $4 = LoadLocal $2 // reference the function
234 - *
235 - * With flag enabled:
236 - * $1 = LoadGlobal useCallback // dead code
237 - * $2 = FunctionExpression ... // the callback being memoized
238 - * $3 = ArrayExpression ... // deps array (dead code)
239 - * .. = Memoize ... // memoize dependencies
240 - * $n = Memoize $2 // reference the function
241 - * $4 = LoadLocal $2 // reference the function
242 - *
243 - * Note that Memoize does not produce a result and is called for its side effects
244 - * only.
245 - */
246 - const functionExpression = functions.get(fn.identifier.id);
247 - if (functionExpression !== undefined) {
248 - for (const operand of eachInstructionValueOperand(
249 - functionExpression
250 - )) {
251 - const temp = createTemporaryPlace(func.env);
252 - nextInstructions.push({
253 - id: makeInstructionId(0),
254 - lvalue: temp,
255 - value: {
256 - kind: "Memoize",
257 - value: { ...operand },
258 - loc: instr.loc,
259 - },
260 - loc: instr.loc,
261 - });
262 - }
263 - }
264 - nextInstructions.push(instr);
287 + loc: fnPlace.loc,
288 + };
289
266 - const temp = createTemporaryPlace(func.env);
267 - nextInstructions.push({
268 - id: makeInstructionId(0),
269 - lvalue: { ...temp },
270 - value: {
271 - kind: "Memoize",
272 - value: {
273 - kind: "Identifier",
274 - identifier: fn.identifier,
275 - effect: Effect.Unknown,
276 - reactive: false,
277 - loc: instr.value.loc,
278 - },
279 - loc: instr.value.loc,
280 - },
281 - loc: instr.loc,
282 - });
283 - continue;
284 - }
285 - }
286 - }
290 + const [startMarker, finishMarker] = makeManualMemoizationMarkers(
291 + fnPlace,
292 + func.env,
293 + // Next PR will replace this with depslist from source
294 + [...eachInstructionValueOperand(inlineMemoFn.value)],
295 + memoDecl
296 + );
297 +
298 + /*
299 + * This PR reorders startMarker to right before the inlineMemoFn
300 + * since startMarker references inlineMemoFn.deps.
301 + * Next PR will move startMarker earlier, to after the `useMemo`/
302 + * `useCallback` load itself (as it also changes startMarker to
303 + * not reference lowered deps anymore).
304 + */
305 + queuedInserts.set(inlineMemoFn.id, {
306 + kind: "before",
307 + value: startMarker,
308 + });
309 + queuedInserts.set(instr.id, { kind: "after", value: finishMarker });
310 + continue;
311 + }
312 + }
313 + } else {
314 + collectTemporaries(instr, func.env, sidemap);
315 + }
316 + }
317 + }
318 +
319 + /**
320 + * Phase 2: Insert manual memoization markers as needed
321 + */
322 + if (queuedInserts.size > 0) {
323 + let hasChanges = false;
324 + for (const [_, block] of func.body.blocks) {
325 + let nextInstructions: Array<Instruction> | null = null;
326 + for (let i = 0; i < block.instructions.length; i++) {
327 + const instr = block.instructions[i];
328 + const insertInstr = queuedInserts.get(instr.id);
329 + if (insertInstr != null) {
330 + nextInstructions = nextInstructions ?? block.instructions.slice(0, i);
331 + if (insertInstr.kind === "before") {
332 + nextInstructions.push(insertInstr.value);
333 + nextInstructions.push(instr);
334 + } else {
335 + nextInstructions.push(instr);
336 + nextInstructions.push(insertInstr.value);
337 }
288 - break;
338 + } else if (nextInstructions != null) {
339 + nextInstructions.push(instr);
340 }
341 }
342 if (nextInstructions !== null) {
292 - nextInstructions.push(instr);
343 + block.instructions = nextInstructions;
344 + hasChanges = true;
345 }
346 }
295 - if (nextInstructions !== null) {
296 - block.instructions = nextInstructions;
297 - hasChanges = true;
347 +
348 + if (hasChanges) {
349 + markInstructionIds(func.body);
350 }
351 }
300 - if (hasChanges) {
301 - markInstructionIds(func.body);
302 - }
352 }
compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts
+18 -15
@@ -1391,21 +1391,24 @@ function inferBlock(
1391 state.alias(lvalue, instrValue.value);
1392 continue;
1393 }
1394 - case "Memoize": {
1395 - if (env.config.enablePreserveExistingMemoizationGuarantees) {
1396 - state.reference(
1397 - instrValue.value,
1398 - functionEffects,
1399 - Effect.Freeze,
1400 - ValueReason.Other
1401 - );
1402 - } else {
1403 - state.reference(
1404 - instrValue.value,
1405 - functionEffects,
1406 - Effect.Read,
1407 - ValueReason.Other
1408 - );
1394 + case "StartMemoize":
1395 + case "FinishMemoize": {
1396 + for (const val of eachInstructionValueOperand(instrValue)) {
1397 + if (env.config.enablePreserveExistingMemoizationGuarantees) {
1398 + state.reference(
1399 + val,
1400 + functionEffects,
1401 + Effect.Freeze,
1402 + ValueReason.Other
1403 + );
1404 + } else {
1405 + state.reference(
1406 + val,
1407 + functionEffects,
1408 + Effect.Read,
1409 + ValueReason.Other
1410 + );
1411 + }
1412 }
1413 const lvalue = instr.lvalue;
1414 lvalue.effect = Effect.ConditionallyMutate;
compiler/packages/babel-plugin-react-forget/src/Optimization/DeadCodeElimination.ts
+2 -1
@@ -338,7 +338,8 @@ function pruneableValue(value: InstructionValue, state: State): boolean {
338 case "StoreContext": {
339 return false;
340 }
341 - case "Memoize": {
341 + case "StartMemoize":
342 + case "FinishMemoize": {
343 /**
344 * This instruction is used by the @enablePreserveExistingMemoizationGuarantees feature
345 * to preserve information about memoization semantics in the original code. We can't
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts
+6 -2
@@ -939,7 +939,10 @@ function codegenInstructionNullable(
939 assertExhaustive(kind, `Unexpected instruction kind '${kind}'`);
940 }
941 }
942 - } else if (instr.value.kind === "Memoize") {
942 + } else if (
943 + instr.value.kind === "StartMemoize" ||
944 + instr.value.kind === "FinishMemoize"
945 + ) {
946 return null;
947 } else if (instr.value.kind === "Debugger") {
948 return t.debuggerStatement();
@@ -1787,7 +1790,8 @@ function codegenInstructionValue(
1790 break;
1791 }
1792 case "ReactiveFunctionValue":
1790 - case "Memoize":
1793 + case "StartMemoize":
1794 + case "FinishMemoize":
1795 case "Debugger":
1796 case "DeclareLocal":
1797 case "DeclareContext":
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/InferReactiveScopeVariables.ts
+2 -1
@@ -154,7 +154,8 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
154 case "NextIterableOf":
155 case "NextPropertyOf":
156 case "Debugger":
157 - case "Memoize":
157 + case "StartMemoize":
158 + case "FinishMemoize":
159 case "UnaryExpression":
160 case "BinaryExpression":
161 case "PropertyLoad": {
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts
+4 -3
@@ -453,7 +453,8 @@ function computeMemoizationInputs(
453 };
454 }
455 case "NextPropertyOf":
456 - case "Memoize":
456 + case "StartMemoize":
457 + case "FinishMemoize":
458 case "Debugger":
459 case "ComputedDelete":
460 case "PropertyDelete":
@@ -926,8 +927,8 @@ class PruneScopesTransform extends ReactiveFunctionTransform<
927 * need to be memoized. Remove associated `Memoize` instructions so that
928 * we don't report false positives on "missing" memoization of these values.
929 */
929 - if (instruction.value.kind === "Memoize") {
930 - const identifier = instruction.value.value.identifier;
930 + if (instruction.value.kind === "FinishMemoize") {
931 + const identifier = instruction.value.decl.identifier;
932 if (
933 identifier.scope !== null &&
934 this.prunedScopes.has(identifier.scope.id)
compiler/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts
+2 -1
@@ -335,7 +335,8 @@ function* generateInstructionTypes(
335 case "NextIterableOf":
336 case "UnsupportedNode":
337 case "Debugger":
338 - case "Memoize": {
338 + case "FinishMemoize":
339 + case "StartMemoize": {
340 break;
341 }
342 default:
compiler/packages/babel-plugin-react-forget/src/Validation/ValidatePreservedManualMemoization.ts
+19 -14
@@ -14,6 +14,7 @@ import {
14 ReactiveScopeBlock,
15 ScopeId,
16 } from "../HIR";
17 +import { eachInstructionValueOperand } from "../HIR/visitors";
18 import { isMutable } from "../ReactiveScopes/InferReactiveScopeVariables";
19 import {
20 ReactiveFunctionVisitor,
@@ -71,20 +72,24 @@ class Visitor extends ReactiveFunctionVisitor<CompilerError> {
72 state: CompilerError
73 ): void {
74 this.traverseInstruction(instruction, state);
74 - if (instruction.value.kind === "Memoize") {
75 - const value = instruction.value.value;
76 - if (
77 - isMutable(instruction as Instruction, value) ||
78 - isUnmemoized(value.identifier, this.scopes)
79 - ) {
80 - state.push({
81 - reason:
82 - "This value was manually memoized, but cannot be memoized under Forget because it may be mutated after it is memoized",
83 - description: null,
84 - severity: ErrorSeverity.InvalidReact,
85 - loc: typeof instruction.loc !== "symbol" ? instruction.loc : null,
86 - suggestions: null,
87 - });
75 + if (
76 + instruction.value.kind === "StartMemoize" ||
77 + instruction.value.kind === "FinishMemoize"
78 + ) {
79 + for (const value of eachInstructionValueOperand(instruction.value)) {
80 + if (
81 + isMutable(instruction as Instruction, value) ||
82 + isUnmemoized(value.identifier, this.scopes)
83 + ) {
84 + state.push({
85 + reason:
86 + "This value was manually memoized, but cannot be memoized under Forget because it may be mutated after it is memoized",
87 + description: null,
88 + severity: ErrorSeverity.InvalidReact,
89 + loc: typeof instruction.loc !== "symbol" ? instruction.loc : null,
90 + suggestions: null,
91 + });
92 + }
93 }
94 }
95 }
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-useCallback-accesses-ref-mutated-later-via-function-preserve-memoization.expect.md
+2 -2
@@ -37,7 +37,7 @@ export const FIXTURE_ENTRYPOINT = {
37 5 | const ref = useRef({ inner: null });
38 6 |
39 > 7 | const onChange = useCallback((event) => {
40 - | ^^^^^^^^^^^^^^^^^^^^^^^^
40 + | ^^^^^^^^^^^^
41 > 8 | // The ref should still be mutable here even though function deps are frozen in
42 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
43 > 9 | // @enablePreserveExistingMemoizationGuarantees mode
@@ -45,7 +45,7 @@ export const FIXTURE_ENTRYPOINT = {
45 > 10 | ref.current.inner = event.target.value;
46 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
47 > 11 | });
48 - | ^^^^^ [ReactForget] InvalidReact: This value was manually memoized, but cannot be memoized under Forget because it may be mutated after it is memoized (7:11)
48 + | ^^^^ [ReactForget] InvalidReact: This value was manually memoized, but cannot be memoized under Forget because it may be mutated after it is memoized (7:11)
49
50 [ReactForget] InvalidReact: This value was manually memoized, but cannot be memoized under Forget because it may be mutated after it is memoized (7:11)
51 12 |
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-useCallback-set-ref-nested-property-ref-modified-later-preserve-memoization.expect.md
+2 -2
@@ -34,7 +34,7 @@ export const FIXTURE_ENTRYPOINT = {
34 5 | const ref = useRef({ inner: null });
35 6 |
36 > 7 | const onChange = useCallback((event) => {
37 - | ^^^^^^^^^^^^^^^^^^^^^^^^
37 + | ^^^^^^^^^^^^
38 > 8 | // The ref should still be mutable here even though function deps are frozen in
39 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
40 > 9 | // @enablePreserveExistingMemoizationGuarantees mode
@@ -42,7 +42,7 @@ export const FIXTURE_ENTRYPOINT = {
42 > 10 | ref.current.inner = event.target.value;
43 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
44 > 11 | });
45 - | ^^^^^ [ReactForget] InvalidReact: This value was manually memoized, but cannot be memoized under Forget because it may be mutated after it is memoized (7:11)
45 + | ^^^^ [ReactForget] InvalidReact: This value was manually memoized, but cannot be memoized under Forget because it may be mutated after it is memoized (7:11)
46
47 [ReactForget] InvalidReact: This value was manually memoized, but cannot be memoized under Forget because it may be mutated after it is memoized (7:11)
48 12 |
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-named-function.expect.md
+20 -5
@@ -2,22 +2,32 @@
2 ## Input
3
4 ```javascript
5 -function Component(props) {
6 - const x = useMemo(someHelper, []);
5 +import { useMemo } from "react";
6 +import { makeArray } from "shared-runtime";
7 +
8 +function Component() {
9 + const x = useMemo(makeArray, []);
10 return x;
11 }
12
13 +export const FIXTURE_ENTRYPOINT = {
14 + fn: Component,
15 + params: [{}],
16 +};
17 +
18 ```
19
20 ## Code
21
22 ```javascript
15 -import { unstable_useMemoCache as useMemoCache } from "react";
16 -function Component(props) {
23 +import { useMemo, unstable_useMemoCache as useMemoCache } from "react";
24 +import { makeArray } from "shared-runtime";
25 +
26 +function Component() {
27 const $ = useMemoCache(1);
28 let t0;
29 if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
20 - t0 = someHelper();
30 + t0 = makeArray();
31 $[0] = t0;
32 } else {
33 t0 = $[0];
@@ -26,5 +36,10 @@ function Component(props) {
36 return x;
37 }
38
39 +export const FIXTURE_ENTRYPOINT = {
40 + fn: Component,
41 + params: [{}],
42 +};
43 +
44 ```
45
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-named-function.js deleted
-4
@@ -1,4 +0,0 @@
1 -function Component(props) {
2 - const x = useMemo(someHelper, []);
3 - return x;
4 -}
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-named-function.ts new
+12
@@ -0,0 +1,12 @@
1 +import { useMemo } from "react";
2 +import { makeArray } from "shared-runtime";
3 +
4 +function Component() {
5 + const x = useMemo(makeArray, []);
6 + return x;
7 +}
8 +
9 +export const FIXTURE_ENTRYPOINT = {
10 + fn: Component,
11 + params: [{}],
12 +};