From 4d7d2cf4dc149deca936912f714493aab085edf8 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Thu, 8 Jun 2023 14:02:04 -0400 Subject: [PATCH] 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. --- .../src/HIR/PrintHIR.ts | 5 +- .../src/SSA/EnterSSA.ts | 59 ++++++++++++++++--- .../src/TypeInference/InferTypes.ts | 17 +++--- .../compiler/array-at-closure.expect.md | 35 +++++------ ...ble-reassigned-outside-of-lambda.expect.md | 0 ...t-variable-reassigned-outside-of-lambda.js | 0 ...ed-scope-declarations-and-locals.expect.md | 4 +- 7 files changed, 85 insertions(+), 35 deletions(-) rename compiler/forget/{ => packages/babel-plugin-react-forget}/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.expect.md (100%) rename compiler/forget/{ => packages/babel-plugin-react-forget}/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.js (100%) diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts b/compiler/forget/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts index 753ae91beb..6a23e532ba 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts +++ b/compiler/forget/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts @@ -431,7 +431,10 @@ export function printInstructionValue(instrValue: ReactiveValue): string { const deps = instrValue.dependencies .map((dep) => printPlace(dep)) .join(","); - value = `Function @deps[${deps}]:\n${fn}`; + const context = instrValue.loweredFunc.context + .map((dep) => printPlace(dep)) + .join(","); + value = `Function @deps[${deps}] @context[${context}]:\n${fn}`; break; } case "TaggedTemplateExpression": { diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts b/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts index 37020eb4dd..2ecfe2e8c3 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts +++ b/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts @@ -46,8 +46,8 @@ class SSABuilder { #unknown: Set = new Set(); #context: Set = new Set(); - constructor(env: Environment, blocks: Map) { - this.#blocks = blocks; + constructor(env: Environment, blocks: ReadonlyMap) { + this.#blocks = new Map(blocks); this.#env = env; } @@ -55,6 +55,18 @@ class SSABuilder { return this.#env.nextIdentifierId; } + defineFunction(func: HIRFunction): void { + for (const [id, block] of func.body.blocks) { + this.#blocks.set(id, block); + } + } + + enter(fn: () => void): void { + const current = this.#current; + fn(); + this.#current = current; + } + state(): State { invariant( this.#current !== null, @@ -125,7 +137,9 @@ class SSABuilder { if (block.preds.size == 0) { // We're at the entry block and haven't found our defintion yet. // console.log( - // `Unable to find "${printIdentifier(oldId)}", assuming it's a global` + // `Unable to find "${printIdentifier( + // oldId + // )}" in bb${blockId}, assuming it's a global` // ); this.#unknown.add(oldId); return oldId; @@ -213,8 +227,16 @@ class SSABuilder { } export default function enterSSA(func: HIRFunction): void { - const visitedBlocks: Set = new Set(); const builder = new SSABuilder(func.env, func.body.blocks); + enterSSAImpl(func, builder, func.body.entry); +} + +function enterSSAImpl( + func: HIRFunction, + builder: SSABuilder, + rootEntry: BlockId +): void { + const visitedBlocks: Set = new Set(); for (const [blockId, block] of func.body.blocks) { invariant( !visitedBlocks.has(block), @@ -224,8 +246,14 @@ export default function enterSSA(func: HIRFunction): void { builder.startBlock(block); - if (func.body.entry === blockId) { - func.context = func.context.map((p) => builder.defineContext(p)); + if (blockId === rootEntry) { + // NOTE: func.context should be empty for the root function + if (func.context.length !== 0) { + CompilerError.invariant( + `Expected function context to be empty for outer function declarations`, + func.loc + ); + } func.params = func.params.map((p) => builder.definePlace(p)); } @@ -234,7 +262,24 @@ export default function enterSSA(func: HIRFunction): void { mapInstructionOperands(instr, (place) => builder.getPlace(place)); if (instr.value.kind === "FunctionExpression") { - enterSSA(instr.value.loweredFunc); + const loweredFunc = instr.value.loweredFunc; + const entry = loweredFunc.body.blocks.get(loweredFunc.body.entry)!; + invariant( + entry.preds.size === 0, + "Expected function expression entry block to have zero predecessors" + ); + entry.preds.add(blockId); + builder.defineFunction(loweredFunc); + builder.enter(() => { + loweredFunc.context = loweredFunc.context.map((p) => + builder.getPlace(p) + ); + loweredFunc.params = loweredFunc.params.map((p) => + builder.definePlace(p) + ); + enterSSAImpl(loweredFunc, builder, rootEntry); + }); + entry.preds.clear(); } } diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts b/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts index 0b8acae0be..0f42adab0b 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts +++ b/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts @@ -65,8 +65,11 @@ function apply(func: HIRFunction, unifier: Unifier): void { for (const place of eachInstructionOperand(instr)) { place.identifier.type = unifier.get(place.identifier.type); } - const { lvalue } = instr; + const { lvalue, value } = instr; lvalue.identifier.type = unifier.get(lvalue.identifier.type); + if (value.kind === "FunctionExpression") { + apply(value.loweredFunc, unifier); + } } } } @@ -125,14 +128,13 @@ function* generateInstructionTypes( break; } - // For now, we won't infer types for context variables - case "StoreContext": { - break; - } + // We intentionally do not infer types for context variables + case "DeclareContext": + case "StoreContext": case "LoadContext": { - yield equation(left, value.place.identifier.type); break; } + case "StoreLocal": { yield equation(left, value.value.identifier.type); yield equation( @@ -245,12 +247,11 @@ function* generateInstructionTypes( } case "FunctionExpression": { - inferTypes(value.loweredFunc); + yield* generate(value.loweredFunc); break; } case "DeclareLocal": - case "DeclareContext": case "NewExpression": case "JsxExpression": case "JsxFragment": diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md index 330045c423..eb63670c41 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md +++ b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md @@ -20,41 +20,42 @@ function Component(props) { import { unstable_useMemoCache as useMemoCache } from "react"; function Component(props) { const $ = useMemoCache(7); - const c_0 = $[0] !== props.x; - let t0; + const t0 = props.x; + const c_0 = $[0] !== t0; + let t1; if (c_0) { - t0 = foo(props.x); - $[0] = props.x; - $[1] = t0; + t1 = foo(t0); + $[0] = t0; + $[1] = t1; } else { - t0 = $[1]; + t1 = $[1]; } - const x = t0; + const x = t1; const c_2 = $[2] !== props; const c_3 = $[3] !== x; - let t1; + let t2; if (c_2 || c_3) { - t1 = function () { + t2 = function () { const arr = [...bar(props)]; return arr.at(x); }; $[2] = props; $[3] = x; - $[4] = t1; + $[4] = t2; } else { - t1 = $[4]; + t2 = $[4]; } - const fn = t1; + const fn = t2; const c_5 = $[5] !== fn; - let t2; + let t3; if (c_5) { - t2 = fn(); + t3 = fn(); $[5] = fn; - $[6] = t2; + $[6] = t3; } else { - t2 = $[6]; + t3 = $[6]; } - const fnResult = t2; + const fnResult = t3; return fnResult; } diff --git a/compiler/forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.expect.md b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.expect.md similarity index 100% rename from compiler/forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.expect.md rename to compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.expect.md diff --git a/compiler/forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.js b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.js similarity index 100% rename from compiler/forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.js rename to compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.js diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md index 431e23cdd7..536652b111 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md +++ b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md @@ -38,8 +38,8 @@ function Component(props) { if (c_0) { const allUrls = []; - const { media: t85, comments, urls } = post; - media = t85; + const { media: t84, comments, urls } = post; + media = t84; const c_3 = $[3] !== comments.length; let t0; if (c_3) {