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

[hir] Wire up object method lowering and codegen

Object methods are lowered to functions and added to ObjectExpression. The codegen is interesting because we shouldn't emit code that lowers the object method into a separate statement and then stores it into an object expression. An shorthand object method has different semantics than an object method using the function syntax, so we need to preserve the shorthand object method syntax in the generated code. To do this, we don't immediately generate an AST node for the ObjectMethod but instead store it in a side table during codegen. Only when emitting code for an ObjectExpression, we lookup this side table and emit the object method inline in the body.

Sathya Gunasekaran committed Oct 2, 2023 at 19:26 UTC f8996476ff348b603eccf5eda49e570984845369
24 files changed +331 -25
compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts
+33
@@ -1277,6 +1277,23 @@ function lowerStatement(
1277 }
1278 }
1279
1280 +function lowerObjectMethod(
1281 + builder: HIRBuilder,
1282 + property: NodePath<t.ObjectMethod>
1283 +): InstructionValue {
1284 + const loc = property.node.loc ?? GeneratedSource;
1285 + const loweredFunc = lowerFunction(builder, property);
1286 + if (!loweredFunc) {
1287 + return { kind: "UnsupportedNode", node: property.node, loc: loc };
1288 + }
1289 +
1290 + return {
1291 + kind: "ObjectMethod",
1292 + loc,
1293 + loweredFunc,
1294 + };
1295 +}
1296 +
1297 function lowerObjectPropertyKey(
1298 builder: HIRBuilder,
1299 key: t.PrivateName | t.Expression
@@ -1379,6 +1396,22 @@ function lowerExpression(
1396 kind: "Spread",
1397 place,
1398 });
1399 + } else if (propertyPath.isObjectMethod()) {
1400 + const method = lowerObjectMethod(builder, propertyPath);
1401 + const place = lowerValueToTemporary(builder, method);
1402 + const loweredKey = lowerObjectPropertyKey(
1403 + builder,
1404 + propertyPath.node.key
1405 + );
1406 + if (!loweredKey) {
1407 + continue;
1408 + }
1409 + properties.push({
1410 + kind: "ObjectProperty",
1411 + type: "method",
1412 + place,
1413 + key: loweredKey,
1414 + });
1415 } else {
1416 builder.errors.push({
1417 reason: `(BuildHIR::lowerExpression) Handle ${propertyPath.type} properties in ObjectExpression`,
compiler/packages/babel-plugin-react-forget/src/HIR/FindContextIdentifiers.ts
+9 -1
@@ -9,6 +9,7 @@ type FindContextIdentifierState = {
9 | NodePath<t.FunctionDeclaration>
10 | NodePath<t.FunctionExpression>
11 | NodePath<t.ArrowFunctionExpression>
12 + | NodePath<t.ObjectMethod>
13 >;
14 reassigned: Set<t.Identifier>;
15 referenced: Set<t.Identifier>;
@@ -75,6 +76,12 @@ export function findContextIdentifiers(
76 const left = path.get("left");
77 handleAssignment(state.reassigned, left);
78 },
79 + ObjectMethod(
80 + fn: NodePath<t.ObjectMethod>,
81 + state: FindContextIdentifierState
82 + ): void {
83 + state.currentLambda.push(fn);
84 + },
85 Identifier(
86 path: NodePath<t.Identifier>,
87 state: FindContextIdentifierState
@@ -93,7 +100,8 @@ function handleIdentifier(
100 currentLambda:
101 | NodePath<t.FunctionDeclaration>
102 | NodePath<t.FunctionExpression>
96 - | NodePath<t.ArrowFunctionExpression>,
103 + | NodePath<t.ArrowFunctionExpression>
104 + | NodePath<t.ObjectMethod>,
105 referenced: Set<t.Identifier>,
106 path: NodePath<t.Identifier>
107 ): void {
compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts
+1
@@ -732,6 +732,7 @@ export type InstructionValue =
732 properties: Array<ObjectProperty | SpreadPattern>;
733 loc: SourceLocation;
734 }
735 + | ObjectMethod
736 | ArrayExpression
737 | { kind: "JsxFragment"; children: Array<Place>; loc: SourceLocation }
738 | {
compiler/packages/babel-plugin-react-forget/src/HIR/MergeConsecutiveBlocks.ts
+4 -1
@@ -32,7 +32,10 @@ export function mergeConsecutiveBlocks(fn: HIRFunction): void {
32 const merged = new MergedBlocks();
33 for (const [, block] of fn.body.blocks) {
34 for (const instr of block.instructions) {
35 - if (instr.value.kind === "FunctionExpression") {
35 + if (
36 + instr.value.kind === "FunctionExpression" ||
37 + instr.value.kind === "ObjectMethod"
38 + ) {
39 mergeConsecutiveBlocks(instr.value.loweredFunc.func);
40 }
41 }
compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts
+1
@@ -462,6 +462,7 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
462 )}]`;
463 break;
464 }
465 + case "ObjectMethod":
466 case "FunctionExpression": {
467 const fn = printFunction(instrValue.loweredFunc.func)
468 .split("\n")
compiler/packages/babel-plugin-react-forget/src/HIR/visitors.ts
+2
@@ -169,6 +169,7 @@ export function* eachInstructionValueOperand(
169 }
170 break;
171 }
172 + case "ObjectMethod":
173 case "FunctionExpression": {
174 yield* instrValue.loweredFunc.dependencies;
175 break;
@@ -456,6 +457,7 @@ export function mapInstructionOperands(
457 instrValue.children = instrValue.children.map((e) => fn(e));
458 break;
459 }
460 + case "ObjectMethod":
461 case "FunctionExpression": {
462 instrValue.loweredFunc.dependencies =
463 instrValue.loweredFunc.dependencies.map((d) => fn(d));
compiler/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts
+7 -6
@@ -8,12 +8,12 @@
8 import { CompilerError } from "../CompilerError";
9 import {
10 Effect,
11 - FunctionExpression,
11 HIRFunction,
12 Identifier,
13 isRefValueType,
14 isSetStateType,
15 isUseRefType,
16 + LoweredFunction,
17 Place,
18 ReactiveScopeDependency,
19 } from "../HIR";
@@ -70,9 +70,10 @@ export default function analyseFunctions(func: HIRFunction): void {
70 for (const [_, block] of func.body.blocks) {
71 for (const instr of block.instructions) {
72 switch (instr.value.kind) {
73 + case "ObjectMethod":
74 case "FunctionExpression": {
75 lower(instr.value.loweredFunc.func);
75 - infer(instr.value, state, func.context);
76 + infer(instr.value.loweredFunc, state, func.context);
77 break;
78 }
79 case "PropertyLoad": {
@@ -113,12 +114,12 @@ function lower(func: HIRFunction): void {
114 }
115
116 function infer(
116 - value: FunctionExpression,
117 + loweredFunc: LoweredFunction,
118 state: IdentifierState,
119 context: Place[]
120 ): void {
121 const mutations = new Map<string, Effect>();
121 - for (const operand of value.loweredFunc.func.context) {
122 + for (const operand of loweredFunc.func.context) {
123 if (
124 isMutatedOrReassigned(operand.identifier) &&
125 operand.identifier.name !== null
@@ -128,7 +129,7 @@ function infer(
129 operand.identifier.mutableRange.end = operand.identifier.mutableRange.start;
130 }
131
131 - for (const dep of value.loweredFunc.dependencies) {
132 + for (const dep of loweredFunc.dependencies) {
133 let name: string | null = null;
134
135 if (state.properties.has(dep.identifier)) {
@@ -174,7 +175,7 @@ function infer(
175 const effect = mutations.get(place.identifier.name);
176 if (effect !== undefined) {
177 place.effect = effect === Effect.Unknown ? Effect.Capture : effect;
177 - value.loweredFunc.dependencies.push(place);
178 + loweredFunc.dependencies.push(place);
179 }
180 }
181 }
compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts
+1
@@ -714,6 +714,7 @@ function inferBlock(
714 valueKind = ValueKind.Immutable;
715 break;
716 }
717 + case "ObjectMethod":
718 case "FunctionExpression": {
719 let hasMutableOperand = false;
720 for (const operand of eachInstructionOperand(instr)) {
compiler/packages/babel-plugin-react-forget/src/Optimization/ConstantPropagation.ts
+4 -1
@@ -106,7 +106,10 @@ function applyConstantPropagation(
106 const functionDependencies = new Set<IdentifierId>();
107 for (const [, block] of fn.body.blocks) {
108 for (const instr of block.instructions) {
109 - if (instr.value.kind === "FunctionExpression") {
109 + if (
110 + instr.value.kind === "FunctionExpression" ||
111 + instr.value.kind === "ObjectMethod"
112 + ) {
113 for (const operand of instr.value.loweredFunc.dependencies) {
114 functionDependencies.add(operand.identifier.id);
115 }
compiler/packages/babel-plugin-react-forget/src/Optimization/DeadCodeElimination.ts
+1
@@ -234,6 +234,7 @@ function pruneableValue(value: InstructionValue, state: State): boolean {
234 case "ArrayExpression":
235 case "BinaryExpression":
236 case "ComputedLoad":
237 + case "ObjectMethod":
238 case "FunctionExpression":
239 case "LoadLocal":
240 case "JsxExpression":
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts
+60 -9
@@ -16,6 +16,7 @@ import {
16 IdentifierId,
17 InstructionKind,
18 JsxAttribute,
19 + ObjectMethod,
20 ObjectPropertyKey,
21 Pattern,
22 Place,
@@ -118,6 +119,8 @@ class Context {
119 #declarations: Set<IdentifierId> = new Set();
120 temp: Temporaries = new Map();
121 errors: CompilerError = new CompilerError();
122 + objectMethods: Map<IdentifierId, ObjectMethod> = new Map();
123 +
124 constructor(env: Environment, fnName: string) {
125 this.env = env;
126 this.fnName = fnName;
@@ -707,6 +710,14 @@ function codegenInstructionNullable(
710 }
711 } else if (instr.value.kind === "Debugger") {
712 return t.debuggerStatement();
713 + } else if (instr.value.kind === "ObjectMethod") {
714 + CompilerError.invariant(instr.lvalue, {
715 + reason: "Expected object methods to have a temp lvalue",
716 + loc: null,
717 + suggestions: null,
718 + });
719 + cx.objectMethods.set(instr.lvalue.identifier.id, instr.value);
720 + return null;
721 } else {
722 const value = codegenInstructionValue(cx, instr.value);
723 const statement = codegenInstruction(cx, instr, value);
@@ -969,15 +980,54 @@ function codegenInstructionValue(
980 for (const property of instrValue.properties) {
981 if (property.kind === "ObjectProperty") {
982 const key = codegenObjectPropertyKey(property.key);
972 - const value = codegenPlace(cx, property.place);
973 - properties.push(
974 - t.objectProperty(
975 - key,
976 - value,
977 - false,
978 - value.type === "Identifier" && value.name === property.key.name
979 - )
980 - );
983 +
984 + switch (property.type) {
985 + case "property": {
986 + const value = codegenPlace(cx, property.place);
987 + properties.push(
988 + t.objectProperty(
989 + key,
990 + value,
991 + false,
992 + value.type === "Identifier" &&
993 + value.name === property.key.name
994 + )
995 + );
996 + break;
997 + }
998 + case "method": {
999 + const method = cx.objectMethods.get(property.place.identifier.id);
1000 + CompilerError.invariant(method, {
1001 + reason: "Expected ObjectMethod instruction",
1002 + loc: null,
1003 + suggestions: null,
1004 + });
1005 + const loweredFunc = method.loweredFunc;
1006 + const reactiveFunction = buildReactiveFunction(loweredFunc.func);
1007 + pruneUnusedLabels(reactiveFunction);
1008 + pruneUnusedLValues(reactiveFunction);
1009 + renameVariables(reactiveFunction);
1010 + const fn = codegenReactiveFunction(reactiveFunction).unwrap();
1011 +
1012 + properties.push(
1013 + t.objectMethod(
1014 + "method",
1015 + key,
1016 + fn.params,
1017 + fn.body,
1018 + false,
1019 + fn.generator,
1020 + fn.async
1021 + )
1022 + );
1023 + break;
1024 + }
1025 + default:
1026 + assertExhaustive(
1027 + property.type,
1028 + `Unexpected property type: ${property.type}`
1029 + );
1030 + }
1031 } else {
1032 properties.push(t.spreadElement(codegenPlace(cx, property.place)));
1033 }
@@ -1296,6 +1346,7 @@ function codegenInstructionValue(
1346 case "DeclareContext":
1347 case "Destructure":
1348 case "StoreLocal":
1349 + case "ObjectMethod":
1350 case "StoreContext": {
1351 CompilerError.invariant(false, {
1352 reason: `Unexpected ${instrValue.kind} in codegenInstructionValue`,
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/InferReactiveScopeVariables.ts
+1
@@ -264,6 +264,7 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
264 case "NewExpression":
265 case "ObjectExpression":
266 case "UnsupportedNode":
267 + case "ObjectMethod":
268 case "FunctionExpression":
269 case "TaggedTemplateExpression": {
270 return true;
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts
+1
@@ -661,6 +661,7 @@ function computeMemoizationInputs(
661 };
662 }
663 case "RegExpLiteral":
664 + case "ObjectMethod":
665 case "FunctionExpression":
666 case "TaggedTemplateExpression":
667 case "ArrayExpression":
compiler/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts
+4 -1
@@ -106,7 +106,10 @@ export function eliminateRedundantPhi(
106 rewritePlace(place, rewrites);
107 }
108
109 - if (instr.value.kind === "FunctionExpression") {
109 + if (
110 + instr.value.kind === "FunctionExpression" ||
111 + instr.value.kind === "ObjectMethod"
112 + ) {
113 const { context } = instr.value.loweredFunc.func;
114 for (const place of context) {
115 rewritePlace(place, rewrites);
compiler/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts
+4 -1
@@ -275,7 +275,10 @@ function enterSSAImpl(
275 mapInstructionOperands(instr, (place) => builder.getPlace(place));
276 mapInstructionLValues(instr, (lvalue) => builder.definePlace(lvalue));
277
278 - if (instr.value.kind === "FunctionExpression") {
278 + if (
279 + instr.value.kind === "FunctionExpression" ||
280 + instr.value.kind === "ObjectMethod"
281 + ) {
282 const loweredFunc = instr.value.loweredFunc.func;
283 const entry = loweredFunc.body.blocks.get(loweredFunc.body.entry)!;
284 CompilerError.invariant(entry.preds.size === 0, {
compiler/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts
+1
@@ -255,6 +255,7 @@ function* generateInstructionTypes(
255 break;
256 }
257
258 + case "ObjectMethod":
259 case "FunctionExpression": {
260 yield* generate(value.loweredFunc.func);
261 break;
compiler/packages/babel-plugin-react-forget/src/Validation/ValidateFrozenLambdas.ts
+9 -3
@@ -15,6 +15,7 @@ import {
15 FunctionExpression,
16 HIRFunction,
17 IdentifierId,
18 + ObjectMethod,
19 Place,
20 isRefValueType,
21 isUseRefType,
@@ -66,6 +67,7 @@ export function validateFrozenLambdas(fn: HIRFunction): void {
67 }
68 for (const instr of block.instructions) {
69 switch (instr.value.kind) {
70 + case "ObjectMethod":
71 case "FunctionExpression": {
72 if (
73 instr.value.loweredFunc.dependencies.some(
@@ -119,7 +121,7 @@ export function validateFrozenLambdas(fn: HIRFunction): void {
121 }
122
123 class State {
122 - lambdas: Map<IdentifierId, FunctionExpression> = new Map();
124 + lambdas: Map<IdentifierId, FunctionExpression | ObjectMethod> = new Map();
125 temporaries: Map<IdentifierId, IdentifierId> = new Map();
126 }
127
@@ -132,9 +134,13 @@ function validateOperand(
134 state.temporaries.get(operand.identifier.id) ?? operand.identifier.id;
135 const lambda = state.lambdas.get(operandId);
136 if (lambda !== undefined) {
135 - // TODO: these seem to always be null, we should try to preserve original names from source
137 + // TODO: these seem to always be null, we should try to preserve original
138 + // names from source
139 + // TODO: figure out how to print object methods as they don't have names
140 const description =
137 - lambda.name !== null && operand.identifier.name !== null
141 + lambda.kind === "FunctionExpression" &&
142 + lambda.name !== null &&
143 + operand.identifier.name !== null
144 ? `\`${lambda.name}\` is a function that may mutate \`${operand.identifier.name}\`. If you must mutate \`${operand.identifier.name}\` try using a React API like useState and use its setter function instead`
145 : null;
146 return new CompilerErrorDetail({
compiler/packages/babel-plugin-react-forget/src/Validation/ValidateNoRefAccesInRender.ts
+1
@@ -59,6 +59,7 @@ export function validateNoRefAccessInRender(fn: HIRFunction): void {
59 }
60 break;
61 }
62 + case "ObjectMethod":
63 case "FunctionExpression": {
64 // functions are allowed to capture refs, so long as the function is not called
65 // during render. see AnalyzeFunctions for how we ensure that functions which
compiler/packages/babel-plugin-react-forget/src/Validation/ValidateNoSetStateInRender.ts
+1
@@ -36,6 +36,7 @@ export function validateNoSetStateInRender(
36 if (unconditionalBlocks.has(block.id)) {
37 for (const instr of block.instructions) {
38 switch (instr.value.kind) {
39 + case "ObjectMethod":
40 case "FunctionExpression": {
41 /**
42 * TODO: setState's return value is considered Frozen, so the lambda's mutable range
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-kitchensink.expect.md
-2
@@ -83,8 +83,6 @@ let moduleLocal = false;
83
84 [ReactForget] Todo: (BuildHIR::lowerStatement) Handle ClassDeclaration statements (5:10)
85
86 -[ReactForget] Todo: (BuildHIR::lowerExpression) Handle ObjectMethod properties in ObjectExpression (12:12)
87 -
86 [ReactForget] Todo: (BuildHIR::lowerStatement) Handle non-variable initialization in ForStatement (20:22)
87
88 [ReactForget] Todo: (BuildHIR::lowerStatement) Handle non-variable initialization in ForStatement (23:25)
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-1.expect.md new
+74
@@ -0,0 +1,74 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +function Component({ a, b }) {
6 + return {
7 + x: function () {
8 + return [a];
9 + },
10 + y() {
11 + return [b];
12 + },
13 + };
14 +}
15 +
16 +export const FIXTURE_ENTRYPOINT = {
17 + fn: Component,
18 + params: [{ x: 1 }, { a: 2 }, { b: 2 }],
19 +};
20 +
21 +```
22 +
23 +## Code
24 +
25 +```javascript
26 +import { unstable_useMemoCache as useMemoCache } from "react";
27 +function Component(t16) {
28 + const $ = useMemoCache(7);
29 + const { a, b } = t16;
30 + const c_0 = $[0] !== a;
31 + let t0;
32 + if (c_0) {
33 + t0 = function () {
34 + return [a];
35 + };
36 + $[0] = a;
37 + $[1] = t0;
38 + } else {
39 + t0 = $[1];
40 + }
41 + const c_2 = $[2] !== b;
42 + let t1;
43 + if (c_2) {
44 + $[2] = b;
45 + $[3] = t1;
46 + } else {
47 + t1 = $[3];
48 + }
49 + const c_4 = $[4] !== t0;
50 + const c_5 = $[5] !== t1;
51 + let t2;
52 + if (c_4 || c_5) {
53 + t2 = {
54 + x: t0,
55 + y() {
56 + return [b];
57 + },
58 + };
59 + $[4] = t0;
60 + $[5] = t1;
61 + $[6] = t2;
62 + } else {
63 + t2 = $[6];
64 + }
65 + return t2;
66 +}
67 +
68 +export const FIXTURE_ENTRYPOINT = {
69 + fn: Component,
70 + params: [{ x: 1 }, { a: 2 }, { b: 2 }],
71 +};
72 +
73 +```
74 +
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-1.js new
+15
@@ -0,0 +1,15 @@
1 +function Component({ a, b }) {
2 + return {
3 + x: function () {
4 + return [a];
5 + },
6 + y() {
7 + return [b];
8 + },
9 + };
10 +}
11 +
12 +export const FIXTURE_ENTRYPOINT = {
13 + fn: Component,
14 + params: [{ x: 1 }, { a: 2 }, { b: 2 }],
15 +};
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-2.expect.md new
+83
@@ -0,0 +1,83 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +function Component({ a, b, c }) {
6 + return {
7 + x: [a],
8 + y() {
9 + return [b];
10 + },
11 + z: { c },
12 + };
13 +}
14 +
15 +export const FIXTURE_ENTRYPOINT = {
16 + fn: Component,
17 + params: [{ x: 1 }, { a: 2 }, { b: 2 }],
18 +};
19 +
20 +```
21 +
22 +## Code
23 +
24 +```javascript
25 +import { unstable_useMemoCache as useMemoCache } from "react";
26 +function Component(t16) {
27 + const $ = useMemoCache(10);
28 + const { a, b, c } = t16;
29 + const c_0 = $[0] !== a;
30 + let t0;
31 + if (c_0) {
32 + t0 = [a];
33 + $[0] = a;
34 + $[1] = t0;
35 + } else {
36 + t0 = $[1];
37 + }
38 + const c_2 = $[2] !== b;
39 + let t1;
40 + if (c_2) {
41 + $[2] = b;
42 + $[3] = t1;
43 + } else {
44 + t1 = $[3];
45 + }
46 + const c_4 = $[4] !== c;
47 + let t2;
48 + if (c_4) {
49 + t2 = { c };
50 + $[4] = c;
51 + $[5] = t2;
52 + } else {
53 + t2 = $[5];
54 + }
55 + const c_6 = $[6] !== t0;
56 + const c_7 = $[7] !== t1;
57 + const c_8 = $[8] !== t2;
58 + let t3;
59 + if (c_6 || c_7 || c_8) {
60 + t3 = {
61 + x: t0,
62 + y() {
63 + return [b];
64 + },
65 + z: t2,
66 + };
67 + $[6] = t0;
68 + $[7] = t1;
69 + $[8] = t2;
70 + $[9] = t3;
71 + } else {
72 + t3 = $[9];
73 + }
74 + return t3;
75 +}
76 +
77 +export const FIXTURE_ENTRYPOINT = {
78 + fn: Component,
79 + params: [{ x: 1 }, { a: 2 }, { b: 2 }],
80 +};
81 +
82 +```
83 +
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-2.js new
+14
@@ -0,0 +1,14 @@
1 +function Component({ a, b, c }) {
2 + return {
3 + x: [a],
4 + y() {
5 + return [b];
6 + },
7 + z: { c },
8 + };
9 +}
10 +
11 +export const FIXTURE_ENTRYPOINT = {
12 + fn: Component,
13 + params: [{ x: 1 }, { a: 2 }, { b: 2 }],
14 +};