[λ] Store and define captured bindings as context

This lets us treat global references differently from captured references.

compiler/forget/src/HIR/BuildHIR.ts

@@ -49,10 +49,21 @@ import HIRBuilder, { Environment } from "./HIRBuilder";
  * grained reactivity.
  */
 export function lower(
-  func: NodePath<t.Function>
+  func: NodePath<t.Function>,
+  capturedRefs?: t.Identifier[]
 ): Result<HIRFunction, CompilerError> {
   const env = new Environment();
   const builder = new HIRBuilder(env);
+  const context: Place[] = [];
+
+  for (const ref of capturedRefs ?? []) {
+    context.push({
+      kind: "Identifier",
+      identifier: builder.resolveBinding(ref),
+      effect: Effect.Unknown,
+      loc: GeneratedSource,
+    });
+  }
 
   // Internal babel is on an older version that does not have hasNode (v7.17)
   // See https://github.com/babel/babel/pull/13940/files for impl
@@ -125,6 +136,7 @@ export function lower(
     id,
     params,
     body: builder.build(),
+    context,
     generator: func.node.generator === true,
     async: func.node.async === true,
     loc: func.node.loc ?? GeneratedSource,
@@ -1300,12 +1312,8 @@ function lowerExpression(
         name = expr.get("id")?.node?.name ?? null;
       }
       const componentScope: Scope = expr.scope.parent.getFunctionParent()!;
-      const dependencies: Array<Place> = gatherCapturedDeps(
-        builder,
-        expr,
-        componentScope
-      );
-      const lowering = lower(expr);
+      const captured = gatherCapturedDeps(builder, expr, componentScope);
+      const lowering = lower(expr, captured.identifiers);
       let loweredFunc: HIRFunction;
       if (lowering.isErr()) {
         lowering
@@ -1340,7 +1348,7 @@ function lowerExpression(
             name,
             params,
             loweredFunc,
-            dependencies,
+            dependencies: captured.refs,
             mutatedDeps: [],
             expr: expr.node,
             loc: exprLoc,
@@ -1919,8 +1927,9 @@ function gatherCapturedDeps(
   builder: HIRBuilder,
   fn: NodePath<t.FunctionExpression | t.ArrowFunctionExpression>,
   componentScope: Scope
-): Array<Place> {
-  const captured: Set<Place> = new Set();
+): { identifiers: t.Identifier[]; refs: Place[] } {
+  const capturedIds: Set<t.Identifier> = new Set();
+  const capturedRefs: Set<Place> = new Set();
 
   // Capture all the scopes from the parent of this function up to and including
   // the component scope.
@@ -1946,9 +1955,10 @@ function gatherCapturedDeps(
       }
 
       path.skip();
-      captured.add(lowerExpressionToPlace(builder, path));
+      capturedIds.add(binding.identifier);
+      capturedRefs.add(lowerExpressionToPlace(builder, path));
     },
   });
 
-  return [...captured];
+  return { identifiers: [...capturedIds], refs: [...capturedRefs] };
 }

compiler/forget/src/HIR/HIR.ts

@@ -184,6 +184,7 @@ export type HIRFunction = {
   id: Identifier | null;
   env: Environment;
   params: Array<Place>;
+  context: Array<Place>;
   body: HIR;
   generator: boolean;
   async: boolean;

compiler/forget/src/HIR/HIRBuilder.ts

@@ -172,6 +172,11 @@ export default class HIRBuilder {
     const node =
       path.scope.getBindingIdentifier(originalName) ??
       getOrAddGlobal(originalName);
+    return this.resolveBinding(node);
+  }
+
+  resolveBinding(node: t.Identifier) {
+    const originalName = node.name;
     let name = originalName;
     let index = 0;
     while (true) {

compiler/forget/src/Inference/AnalyseFunctions.ts

@@ -1,11 +1,4 @@
-import {
-  Effect,
-  HIRFunction,
-  Identifier,
-  mergeConsecutiveBlocks,
-  Place,
-} from "../HIR";
-import { eachInstructionOperand } from "../HIR/visitors";
+import { HIRFunction, Identifier, mergeConsecutiveBlocks, Place } from "../HIR";
 import { constantPropagation } from "../Optimization";
 import { eliminateRedundantPhi, enterSSA } from "../SSA";
 import { inferTypes } from "../TypeInference";
@@ -116,18 +109,13 @@ function analyzeMutatedPlaces(func: HIRFunction): Array<Place> {
       ) {
         mutations.push(...analyzeMutatedPlaces(instr.value.loweredFunc));
       }
-
-      for (const operand of eachInstructionOperand(instr)) {
-        if (isMutated(operand)) {
-          mutations.push(operand);
-        }
-      }
     }
   }
 
+  mutations.push(...func.context.filter((dep) => isMutated(dep.identifier)));
   return mutations;
 }
 
-function isMutated(place: Place): boolean {
-  return place.effect === Effect.Mutate || place.effect === Effect.Store;
+function isMutated(id: Identifier) {
+  return id.mutableRange.end - id.mutableRange.start > 1;
 }

compiler/forget/src/Inference/InferReferenceEffects.ts

@@ -93,6 +93,17 @@ export default function inferReferenceEffects(fn: HIRFunction) {
     initialEnvironment.define(id, value);
   }
 
+  for (const ref of fn.context) {
+    // TODO(gsn): This is a hack.
+    const value: InstructionValue = {
+      kind: "ObjectExpression",
+      properties: null,
+      loc: ref.loc,
+    };
+    initialEnvironment.initialize(value, ValueKind.Mutable);
+    initialEnvironment.define(ref, value);
+  }
+
   for (const param of fn.params) {
     const value: InstructionValue = {
       kind: "Primitive",

compiler/forget/src/SSA/EnterSSA.ts

@@ -195,6 +195,7 @@ export default function enterSSA(func: HIRFunction) {
     builder.startBlock(block);
 
     if (func.body.entry === blockId) {
+      func.context = func.context.map((p) => builder.definePlace(p));
       func.params = func.params.map((p) => builder.definePlace(p));
     }
 

compiler/forget/src/__tests__/fixtures/hir/capture-param-mutate.expect.md

@@ -0,0 +1,97 @@
+
+## Input
+
+```javascript
+function getNativeLogFunction(level) {
+  return function () {
+    let str;
+    if (arguments.length === 1 && typeof arguments[0] === "string") {
+      str = arguments[0];
+    } else {
+      str = Array.prototype.map
+        .call(arguments, function (arg) {
+          return inspect(arg, {
+            depth: 10,
+          });
+        })
+        .join(", ");
+    }
+    const firstArg = arguments[0];
+    let logLevel = level;
+    if (
+      typeof firstArg === "string" &&
+      firstArg.slice(0, 9) === "Warning: " &&
+      logLevel >= LOG_LEVELS.error
+    ) {
+      logLevel = LOG_LEVELS.warn;
+    }
+    if (global.__inspectorLog) {
+      global.__inspectorLog(
+        INSPECTOR_LEVELS[logLevel],
+        str,
+        [].slice.call(arguments),
+        INSPECTOR_FRAMES_TO_SKIP
+      );
+    }
+    if (groupStack.length) {
+      str = groupFormat("", str);
+    }
+    global.nativeLoggingHook(str, logLevel);
+  };
+}
+
+```
+
+## Code
+
+```javascript
+function getNativeLogFunction(level) {
+  const $ = React.useMemoCache();
+  const c_0 = $[0] !== level;
+  let t1;
+  if (c_0) {
+    t1 = function () {
+      let str;
+      if (arguments.length === 1 && typeof arguments[0] === "string") {
+        str = arguments[0];
+      } else {
+        str = Array.prototype.map
+          .call(arguments, function (arg) {
+            return inspect(arg, {
+              depth: 10,
+            });
+          })
+          .join(", ");
+      }
+      const firstArg = arguments[0];
+      let logLevel = level;
+      if (
+        typeof firstArg === "string" &&
+        firstArg.slice(0, 9) === "Warning: " &&
+        logLevel >= LOG_LEVELS.error
+      ) {
+        logLevel = LOG_LEVELS.warn;
+      }
+      if (global.__inspectorLog) {
+        global.__inspectorLog(
+          INSPECTOR_LEVELS[logLevel],
+          str,
+          [].slice.call(arguments),
+          INSPECTOR_FRAMES_TO_SKIP
+        );
+      }
+      if (groupStack.length) {
+        str = groupFormat("", str);
+      }
+      global.nativeLoggingHook(str, logLevel);
+    };
+    $[0] = level;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
+}
+
+```
+      

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

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

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

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

compiler/forget/src/__tests__/fixtures/hir/error._bug_fbsource_pokes_infer_reference_effects.expect.md

@@ -1,52 +0,0 @@
-
-## Input
-
-```javascript
-function getNativeLogFunction(level) {
-  return function () {
-    let str;
-    if (arguments.length === 1 && typeof arguments[0] === "string") {
-      str = arguments[0];
-    } else {
-      str = Array.prototype.map
-        .call(arguments, function (arg) {
-          return inspect(arg, {
-            depth: 10,
-          });
-        })
-        .join(", ");
-    }
-    const firstArg = arguments[0];
-    let logLevel = level;
-    if (
-      typeof firstArg === "string" &&
-      firstArg.slice(0, 9) === "Warning: " &&
-      logLevel >= LOG_LEVELS.error
-    ) {
-      logLevel = LOG_LEVELS.warn;
-    }
-    if (global.__inspectorLog) {
-      global.__inspectorLog(
-        INSPECTOR_LEVELS[logLevel],
-        str,
-        [].slice.call(arguments),
-        INSPECTOR_FRAMES_TO_SKIP
-      );
-    }
-    if (groupStack.length) {
-      str = groupFormat("", str);
-    }
-    global.nativeLoggingHook(str, logLevel);
-  };
-}
-
-```
-
-
-## Error
-
-```
-[ReactForget] Invariant: InferReferenceEffects::kind: Expected at least one value at '<unknown> logLevel$85:TPrimitive' (20:20)
-```
-          
-