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.

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": {

compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts

@@ -46,8 +46,8 @@ class SSABuilder {
   #unknown: Set<Identifier> = new Set();
   #context: Set<Identifier> = new Set();
 
-  constructor(env: Environment, blocks: Map<BlockId, BasicBlock>) {
-    this.#blocks = blocks;
+  constructor(env: Environment, blocks: ReadonlyMap<BlockId, BasicBlock>) {
+    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<BasicBlock> = 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<BasicBlock> = 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();
       }
     }
 

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":

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;
 }
 

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) {