Type inference across function expressions boundaries
Updates InferTypes to perform type inference across function boundaries. Specifically InferTypes is now responsible for driving type inference of function expressions (rather than deferring to AnalyzeFunctions to infer functions), and type inference now traverses into function expressions and infers types of free variables taking into account information from the outer context. This relies on the fact that identifier ids are consistent across function expression boundaries and that all free variables in functions are guaranteed to be effectively `const`, since we promote non-const variables used in function expressions to context variables.
Joe Savona committed
Jun 8, 2023 at 14:02 UTC
4d7d2cf4dc149deca936912f714493aab085edf8
7 files changed
+85
-35
compiler/forget/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts
+4
-1
@@ -431,7 +431,10 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
431
const deps = instrValue.dependencies
432
.map((dep) => printPlace(dep))
433
.join(",");
434
- value = `Function @deps[${deps}]:\n${fn}`;
434
+ const context = instrValue.loweredFunc.context
435
+ .map((dep) => printPlace(dep))
436
+ .join(",");
437
+ value = `Function @deps[${deps}] @context[${context}]:\n${fn}`;
438
break;
439
}
440
case "TaggedTemplateExpression": {
compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts
+52
-7
@@ -46,8 +46,8 @@ class SSABuilder {
46
#unknown: Set<Identifier> = new Set();
47
#context: Set<Identifier> = new Set();
48
49
- constructor(env: Environment, blocks: Map<BlockId, BasicBlock>) {
50
- this.#blocks = blocks;
49
+ constructor(env: Environment, blocks: ReadonlyMap<BlockId, BasicBlock>) {
50
+ this.#blocks = new Map(blocks);
51
this.#env = env;
52
}
53
@@ -55,6 +55,18 @@ class SSABuilder {
55
return this.#env.nextIdentifierId;
56
}
57
58
+ defineFunction(func: HIRFunction): void {
59
+ for (const [id, block] of func.body.blocks) {
60
+ this.#blocks.set(id, block);
61
+ }
62
+ }
63
+
64
+ enter(fn: () => void): void {
65
+ const current = this.#current;
66
+ fn();
67
+ this.#current = current;
68
+ }
69
+
70
state(): State {
71
invariant(
72
this.#current !== null,
@@ -125,7 +137,9 @@ class SSABuilder {
137
if (block.preds.size == 0) {
138
// We're at the entry block and haven't found our defintion yet.
139
// console.log(
128
- // `Unable to find "${printIdentifier(oldId)}", assuming it's a global`
140
+ // `Unable to find "${printIdentifier(
141
+ // oldId
142
+ // )}" in bb${blockId}, assuming it's a global`
143
// );
144
this.#unknown.add(oldId);
145
return oldId;
@@ -213,8 +227,16 @@ class SSABuilder {
227
}
228
229
export default function enterSSA(func: HIRFunction): void {
216
- const visitedBlocks: Set<BasicBlock> = new Set();
230
const builder = new SSABuilder(func.env, func.body.blocks);
231
+ enterSSAImpl(func, builder, func.body.entry);
232
+}
233
+
234
+function enterSSAImpl(
235
+ func: HIRFunction,
236
+ builder: SSABuilder,
237
+ rootEntry: BlockId
238
+): void {
239
+ const visitedBlocks: Set<BasicBlock> = new Set();
240
for (const [blockId, block] of func.body.blocks) {
241
invariant(
242
!visitedBlocks.has(block),
@@ -224,8 +246,14 @@ export default function enterSSA(func: HIRFunction): void {
246
247
builder.startBlock(block);
248
227
- if (func.body.entry === blockId) {
228
- func.context = func.context.map((p) => builder.defineContext(p));
249
+ if (blockId === rootEntry) {
250
+ // NOTE: func.context should be empty for the root function
251
+ if (func.context.length !== 0) {
252
+ CompilerError.invariant(
253
+ `Expected function context to be empty for outer function declarations`,
254
+ func.loc
255
+ );
256
+ }
257
func.params = func.params.map((p) => builder.definePlace(p));
258
}
259
@@ -234,7 +262,24 @@ export default function enterSSA(func: HIRFunction): void {
262
mapInstructionOperands(instr, (place) => builder.getPlace(place));
263
264
if (instr.value.kind === "FunctionExpression") {
237
- enterSSA(instr.value.loweredFunc);
265
+ const loweredFunc = instr.value.loweredFunc;
266
+ const entry = loweredFunc.body.blocks.get(loweredFunc.body.entry)!;
267
+ invariant(
268
+ entry.preds.size === 0,
269
+ "Expected function expression entry block to have zero predecessors"
270
+ );
271
+ entry.preds.add(blockId);
272
+ builder.defineFunction(loweredFunc);
273
+ builder.enter(() => {
274
+ loweredFunc.context = loweredFunc.context.map((p) =>
275
+ builder.getPlace(p)
276
+ );
277
+ loweredFunc.params = loweredFunc.params.map((p) =>
278
+ builder.definePlace(p)
279
+ );
280
+ enterSSAImpl(loweredFunc, builder, rootEntry);
281
+ });
282
+ entry.preds.clear();
283
}
284
}
285
compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts
+9
-8
@@ -65,8 +65,11 @@ function apply(func: HIRFunction, unifier: Unifier): void {
65
for (const place of eachInstructionOperand(instr)) {
66
place.identifier.type = unifier.get(place.identifier.type);
67
}
68
- const { lvalue } = instr;
68
+ const { lvalue, value } = instr;
69
lvalue.identifier.type = unifier.get(lvalue.identifier.type);
70
+ if (value.kind === "FunctionExpression") {
71
+ apply(value.loweredFunc, unifier);
72
+ }
73
}
74
}
75
}
@@ -125,14 +128,13 @@ function* generateInstructionTypes(
128
break;
129
}
130
128
- // For now, we won't infer types for context variables
129
- case "StoreContext": {
130
- break;
131
- }
131
+ // We intentionally do not infer types for context variables
132
+ case "DeclareContext":
133
+ case "StoreContext":
134
case "LoadContext": {
133
- yield equation(left, value.place.identifier.type);
135
break;
136
}
137
+
138
case "StoreLocal": {
139
yield equation(left, value.value.identifier.type);
140
yield equation(
@@ -245,12 +247,11 @@ function* generateInstructionTypes(
247
}
248
249
case "FunctionExpression": {
248
- inferTypes(value.loweredFunc);
250
+ yield* generate(value.loweredFunc);
251
break;
252
}
253
254
case "DeclareLocal":
253
- case "DeclareContext":
255
case "NewExpression":
256
case "JsxExpression":
257
case "JsxFragment":
compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md
+18
-17
@@ -20,41 +20,42 @@ function Component(props) {
20
import { unstable_useMemoCache as useMemoCache } from "react";
21
function Component(props) {
22
const $ = useMemoCache(7);
23
- const c_0 = $[0] !== props.x;
24
- let t0;
23
+ const t0 = props.x;
24
+ const c_0 = $[0] !== t0;
25
+ let t1;
26
if (c_0) {
26
- t0 = foo(props.x);
27
- $[0] = props.x;
28
- $[1] = t0;
27
+ t1 = foo(t0);
28
+ $[0] = t0;
29
+ $[1] = t1;
30
} else {
30
- t0 = $[1];
31
+ t1 = $[1];
32
}
32
- const x = t0;
33
+ const x = t1;
34
const c_2 = $[2] !== props;
35
const c_3 = $[3] !== x;
35
- let t1;
36
+ let t2;
37
if (c_2 || c_3) {
37
- t1 = function () {
38
+ t2 = function () {
39
const arr = [...bar(props)];
40
return arr.at(x);
41
};
42
$[2] = props;
43
$[3] = x;
43
- $[4] = t1;
44
+ $[4] = t2;
45
} else {
45
- t1 = $[4];
46
+ t2 = $[4];
47
}
47
- const fn = t1;
48
+ const fn = t2;
49
const c_5 = $[5] !== fn;
49
- let t2;
50
+ let t3;
51
if (c_5) {
51
- t2 = fn();
52
+ t3 = fn();
53
$[5] = fn;
53
- $[6] = t2;
54
+ $[6] = t3;
55
} else {
55
- t2 = $[6];
56
+ t3 = $[6];
57
}
57
- const fnResult = t2;
58
+ const fnResult = t3;
59
return fnResult;
60
}
61
compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.expect.md
renamed
compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.js
renamed
compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md
+2
-2
@@ -38,8 +38,8 @@ function Component(props) {
38
if (c_0) {
39
const allUrls = [];
40
41
- const { media: t85, comments, urls } = post;
42
- media = t85;
41
+ const { media: t84, comments, urls } = post;
42
+ media = t84;
43
const c_3 = $[3] !== comments.length;
44
let t0;
45
if (c_3) {