[hir] Lower function expressions into HIR

compiler/forget/src/CompilerPipeline.ts

@@ -13,7 +13,11 @@ import {
   mergeConsecutiveBlocks,
   ReactiveFunction,
 } from "./HIR";
-import { inferMutableRanges, inferReferenceEffects } from "./Inference";
+import {
+  analyseFunctions,
+  inferMutableRanges,
+  inferReferenceEffects,
+} from "./Inference";
 import { constantPropagation } from "./Optimization";
 import {
   alignReactiveScopesToBlockScopes,
@@ -65,6 +69,9 @@ export function* run(
   inferTypes(hir);
   yield log({ kind: "hir", name: "InferTypes", value: hir });
 
+  analyseFunctions(hir);
+  yield log({ kind: "hir", name: "analyseFunctions", value: hir });
+
   inferReferenceEffects(hir);
   yield log({ kind: "hir", name: "InferReferenceEffects", value: hir });
 
@@ -163,7 +170,7 @@ export function compile(func: NodePath<t.FunctionDeclaration>): t.Function {
   }
 }
 
-function log(value: CompilerPipelineValue): CompilerPipelineValue {
+export function log(value: CompilerPipelineValue): CompilerPipelineValue {
   switch (value.kind) {
     case "ast": {
       break;

compiler/forget/src/HIR/BuildHIR.ts

@@ -1216,6 +1216,18 @@ function lowerExpression(
         componentScope
       );
       const body = expr.get("body").node;
+      const lowering = lower(expr);
+      let loweredFunc: HIRFunction;
+      if (lowering.isErr()) {
+        lowering.unwrapErr().forEach((e) => builder.pushError(e));
+        return {
+          kind: "OtherStatement",
+          node: expr.node,
+          loc: exprLoc,
+        };
+      }
+      loweredFunc = lowering.unwrap();
+
       const params: Array<string> = expr.get("params").map((p) => {
         todoInvariant(p.isIdentifier(), "handle non identifier params");
         return p.node.name;
@@ -1225,7 +1237,9 @@ function lowerExpression(
         name,
         body,
         params,
+        loweredFunc,
         dependencies,
+        mutatedDeps: [],
         loc: exprLoc,
       };
     }

compiler/forget/src/HIR/HIR.ts

@@ -403,6 +403,10 @@ export type InstructionData =
       name: string | null;
       params: Array<string>;
       dependencies: Array<Place>;
+      // TODO(gsn): Remove this mutatedDeps array and use dependencies as single
+      // source of truth.
+      mutatedDeps: Array<Place>;
+      loweredFunc: HIRFunction;
       body: t.BlockStatement;
     }
 
@@ -413,7 +417,11 @@ export type InstructionData =
    */
   | {
       kind: "OtherStatement";
-      node: t.Statement | t.JSXSpreadChild | t.JSXFragment;
+      node:
+        | t.Statement
+        | t.JSXSpreadChild
+        | t.JSXFragment
+        | t.FunctionExpression;
     };
 
 /**

compiler/forget/src/Inference/AnalyseFunctions.ts

@@ -0,0 +1,128 @@
+import {
+  Effect,
+  HIRFunction,
+  Identifier,
+  mergeConsecutiveBlocks,
+  Place,
+} from "../HIR";
+import { eachInstructionOperand } from "../HIR/visitors";
+import { constantPropagation } from "../Optimization";
+import { eliminateRedundantPhi, enterSSA } from "../SSA";
+import { inferTypes } from "../TypeInference";
+import { logHIRFunction } from "../Utils/logger";
+import { inferMutableRanges } from "./InferMutableRanges";
+import inferReferenceEffects from "./InferReferenceEffects";
+
+type Dependency = {
+  place: Place;
+  path: Array<string> | null;
+};
+
+function declareProperty(
+  properties: Map<Identifier, Dependency>,
+  lvalue: Place,
+  object: Place,
+  property: string
+): void {
+  const objectDependency = properties.get(object.identifier);
+  let nextDependency: Dependency;
+  if (objectDependency === undefined) {
+    nextDependency = { place: object, path: [property] };
+  } else {
+    nextDependency = {
+      place: objectDependency.place,
+      path: [...(objectDependency.path ?? []), property],
+    };
+  }
+  properties.set(lvalue.identifier, nextDependency);
+}
+
+export default function (func: HIRFunction) {
+  const properties: Map<Identifier, Dependency> = new Map();
+
+  for (const [_, block] of func.body.blocks) {
+    for (const instr of block.instructions) {
+      switch (instr.value.kind) {
+        case "FunctionExpression": {
+          instr.value.mutatedDeps = buildMutatedDeps(
+            analyzeMutatedPlaces(instr.value.loweredFunc),
+            instr.value.dependencies,
+            properties
+          );
+          break;
+        }
+        case "PropertyLoad": {
+          declareProperty(
+            properties,
+            instr.lvalue.place,
+            instr.value.object,
+            instr.value.property
+          );
+        }
+      }
+    }
+  }
+}
+
+function buildMutatedDeps(
+  mutations: Place[],
+  capturedDeps: Place[],
+  properties: Map<Identifier, Dependency>
+): Place[] {
+  const mutatedIds: Set<string> = new Set(
+    mutations
+      .map((m) => m.identifier.name)
+      .filter((m) => m !== null) as string[]
+  );
+  const mutatedDeps: Place[] = [];
+
+  for (const dep of capturedDeps) {
+    if (properties.has(dep.identifier)) {
+      let captured = properties.get(dep.identifier)!;
+      let name = captured.place.identifier.name;
+
+      if (name === null || !mutatedIds.has(name)) {
+        continue;
+      }
+
+      mutatedDeps.push(dep);
+    }
+  }
+
+  return mutatedDeps;
+}
+
+function analyzeMutatedPlaces(func: HIRFunction): Array<Place> {
+  mergeConsecutiveBlocks(func);
+  enterSSA(func);
+  eliminateRedundantPhi(func);
+  constantPropagation(func);
+  inferTypes(func);
+  inferReferenceEffects(func);
+  inferMutableRanges(func);
+  logHIRFunction("AnalyseFunction (inner)", func);
+
+  const mutations: Array<Place> = [];
+  for (const [_, block] of func.body.blocks) {
+    for (const instr of block.instructions) {
+      if (
+        instr.value.kind === "FunctionExpression" &&
+        instr.value.loweredFunc !== null
+      ) {
+        mutations.push(...analyzeMutatedPlaces(instr.value.loweredFunc));
+      }
+
+      for (const operand of eachInstructionOperand(instr)) {
+        if (isMutated(operand)) {
+          mutations.push(operand);
+        }
+      }
+    }
+  }
+
+  return mutations;
+}
+
+function isMutated(place: Place): boolean {
+  return place.effect === Effect.Mutate || place.effect === Effect.Store;
+}

compiler/forget/src/Inference/InferAliasForStores.ts

@@ -46,6 +46,11 @@ export function inferAliasForStores(
           maybeAlias(aliases, value.object, value.value, instr.id);
           break;
         }
+        case "FunctionExpression": {
+          for (const dep of value.mutatedDeps) {
+            maybeAlias(aliases, lvalue.place, dep, instr.id);
+          }
+        }
       }
     }
   }

compiler/forget/src/Inference/InferReferenceEffects.ts

@@ -211,11 +211,10 @@ class Environment {
    */
   alias(place: Place, value: Place) {
     const values = this.#variables.get(value.identifier.id);
-    invariant(
-      values != null,
-      "Expected value for identifier `%s` to be initialized.",
-      value.identifier.id
-    );
+    // A value can be undefined if it has been captured from outside scope.
+    if (value === undefined) {
+      return;
+    }
     this.#variables.set(place.identifier.id, new Set(values));
   }
 

compiler/forget/src/Inference/index.ts

@@ -6,4 +6,5 @@
  */
 
 export { inferMutableRanges } from "./InferMutableRanges";
+export { default as analyseFunctions } from "./AnalyseFunctions";
 export { default as inferReferenceEffects } from "./InferReferenceEffects";

compiler/forget/src/__tests__/fixtures/hir/capturing-func-mutate-2.expect.md

@@ -0,0 +1,52 @@
+
+## Input
+
+```javascript
+function component(a, b) {
+  let y = { b };
+  let z = { a };
+  let x = function () {
+    z.a = 2;
+    y.b;
+  };
+  x();
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+function component(a, b) {
+  const $ = React.useMemoCache();
+  const c_0 = $[0] !== b;
+  let y;
+  if (c_0) {
+    y = { b: b };
+    $[0] = b;
+    $[1] = y;
+  } else {
+    y = $[1];
+  }
+  const c_2 = $[2] !== a;
+  const c_3 = $[3] !== y.b;
+  let x;
+  if (c_2 || c_3) {
+    const z = { a: a };
+    x = function () {
+      z.a = 2;
+      y.b;
+    };
+    x();
+    $[2] = a;
+    $[3] = y.b;
+    $[4] = x;
+  } else {
+    x = $[4];
+  }
+  return x;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/hir/capturing-func-mutate-2.js

@@ -0,0 +1,10 @@
+function component(a, b) {
+  let y = { b };
+  let z = { a };
+  let x = function () {
+    z.a = 2;
+    y.b;
+  };
+  x();
+  return x;
+}

compiler/forget/src/__tests__/fixtures/hir/capturing-func-mutate-3.expect.md

@@ -0,0 +1,58 @@
+
+## Input
+
+```javascript
+function component(a, b) {
+  let y = { b };
+  let z = { a };
+  let x = function () {
+    z.a = 2;
+    y.b;
+  };
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+function component(a, b) {
+  const $ = React.useMemoCache();
+  const c_0 = $[0] !== b;
+  let y;
+  if (c_0) {
+    y = { b: b };
+    $[0] = b;
+    $[1] = y;
+  } else {
+    y = $[1];
+  }
+  const c_2 = $[2] !== a;
+  let z;
+  if (c_2) {
+    z = { a: a };
+    $[2] = a;
+    $[3] = z;
+  } else {
+    z = $[3];
+  }
+  const c_4 = $[4] !== z.a;
+  const c_5 = $[5] !== y.b;
+  let x;
+  if (c_4 || c_5) {
+    x = function () {
+      z.a = 2;
+      y.b;
+    };
+    $[4] = z.a;
+    $[5] = y.b;
+    $[6] = x;
+  } else {
+    x = $[6];
+  }
+  return x;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/hir/capturing-func-mutate-3.js

@@ -0,0 +1,9 @@
+function component(a, b) {
+  let y = { b };
+  let z = { a };
+  let x = function () {
+    z.a = 2;
+    y.b;
+  };
+  return x;
+}

compiler/forget/src/__tests__/fixtures/hir/capturing-func-mutate-nested.js

@@ -0,0 +1,10 @@
+// @skip
+// TODO(gsn): This doesn't seem to work correctly. Need to debug more.
+function component(a) {
+  let y = { b: { a } };
+  let x = function () {
+    y.b.a = 2;
+  };
+  x();
+  return x;
+}

compiler/forget/src/__tests__/fixtures/hir/capturing-func-mutate.expect.md

@@ -0,0 +1,52 @@
+
+## Input
+
+```javascript
+function component(a, b) {
+  let z = { a };
+  let y = { b };
+  let x = function () {
+    z.a = 2;
+    y.b;
+  };
+  x();
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+function component(a, b) {
+  const $ = React.useMemoCache();
+  const c_0 = $[0] !== a;
+  const c_1 = $[1] !== b;
+  let x;
+  if (c_0 || c_1) {
+    const z = { a: a };
+    const c_3 = $[3] !== b;
+    let y;
+    if (c_3) {
+      y = { b: b };
+      $[3] = b;
+      $[4] = y;
+    } else {
+      y = $[4];
+    }
+    x = function () {
+      z.a = 2;
+      y.b;
+    };
+    x();
+    $[0] = a;
+    $[1] = b;
+    $[2] = x;
+  } else {
+    x = $[2];
+  }
+  return x;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/hir/capturing-func-mutate.js

@@ -0,0 +1,10 @@
+function component(a, b) {
+  let z = { a };
+  let y = { b };
+  let x = function () {
+    z.a = 2;
+    y.b;
+  };
+  x();
+  return x;
+}

compiler/forget/src/__tests__/fixtures/hir/capturing-var-in-nested-func-decl.js

@@ -1,3 +1,4 @@
+// @skip -- TODO: support lowering Function Declaration in HIR
 function component(a) {
   let z = { a };
   let x = function () {