[hir] Gather dependencies for lambdas

compiler/forget/src/HIR/BuildHIR.ts

@@ -5,7 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import { NodePath } from "@babel/traverse";
+import { NodePath, Scope } from "@babel/traverse";
 import * as t from "@babel/types";
 import { invariant } from "../Utils/CompilerError";
 import todo, { todoInvariant } from "../Utils/todo";
@@ -1151,6 +1151,8 @@ function lowerExpression(
     case "FunctionExpression": {
       const expr = exprPath as NodePath<t.FunctionExpression>;
       const name: string | null = expr.get("id")?.node?.name ?? null;
+      const componentScope: Scope = expr.scope.parent.getFunctionParent()!;
+      const dependencies = gatherCapturedDeps(expr, componentScope);
       const body = expr.get("body").node;
       const params: Array<string> = expr.get("params").map((p) => {
         todoInvariant(p.isIdentifier(), "handle non identifier params");
@@ -1161,6 +1163,7 @@ function lowerExpression(
         name,
         body,
         params,
+        dependencies,
         loc: exprLoc,
       };
     }
@@ -1540,3 +1543,53 @@ function lowerAssignment(
     }
   }
 }
+
+function capturePureScopes(
+  currentScope: Scope,
+  componentScope: Scope
+): Set<Scope> {
+  let pureScopes: Set<Scope> = new Set();
+  while (currentScope) {
+    pureScopes.add(currentScope);
+
+    if (currentScope === componentScope) {
+      break;
+    }
+
+    currentScope = currentScope.parent;
+  }
+  return pureScopes;
+}
+
+function gatherCapturedDeps(
+  fn: NodePath<t.FunctionExpression>,
+  componentScope: Scope
+): Set<t.Identifier> {
+  const captured: Set<t.Identifier> = new Set();
+
+  // Capture all the scopes from the parent of this function up to and including
+  // the component scope.
+  const pureScopes: Set<Scope> = capturePureScopes(
+    fn.scope.parent,
+    componentScope
+  );
+
+  fn.get("body").traverse({
+    Expression(path) {
+      // TODO(gsn): Handle member expressions
+      if (!path.isIdentifier) {
+        return;
+      }
+
+      const id = path as NodePath<t.Identifier>;
+      const binding = id.scope.getBinding(id.node.name);
+      if (binding === undefined || !pureScopes.has(binding.scope)) {
+        return;
+      }
+
+      captured.add(binding.identifier);
+    },
+  });
+
+  return captured;
+}

compiler/forget/src/HIR/HIR.ts

@@ -313,6 +313,7 @@ export type InstructionData =
       kind: "FunctionExpression";
       name: string | null;
       params: Array<string>;
+      dependencies: Set<t.Identifier>;
       body: t.BlockStatement;
     }
 

compiler/forget/src/HIR/PrintHIR.ts

@@ -293,7 +293,10 @@ export function printInstructionValue(instrValue: InstructionValue): string {
     case "FunctionExpression": {
       const params = instrValue.params.join(",");
       const body = generate(instrValue.body).code;
-      value = `Function ${instrValue.name}(${params}){${body}}`;
+      const deps = [...instrValue.dependencies].map((i) => i.name).join(",");
+      value = `Function ${instrValue.name ?? ""} @deps[${deps}] (${
+        params ?? ""
+      }){${body}}`;
       break;
     }
     default: {

compiler/forget/src/__tests__/fixtures/hir/capturing-function-1.expect.md

@@ -0,0 +1,36 @@
+
+## Input
+
+```javascript
+function component() {
+  let z = 100;
+  let x = function () {
+    z;
+  };
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+function component() {
+  const $ = React.useMemoCache();
+  const z = 100;
+  let x;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    x = function () {
+      z;
+    };
+
+    $[0] = x;
+  } else {
+    x = $[0];
+  }
+
+  return x;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/hir/capturing-function-1.js

@@ -0,0 +1,7 @@
+function component() {
+  let z = 100;
+  let x = function () {
+    z;
+  };
+  return x;
+}

compiler/forget/src/__tests__/fixtures/hir/capturing-function-within-block.expect.md

@@ -0,0 +1,40 @@
+
+## Input
+
+```javascript
+function component() {
+  let z = 100;
+  let x;
+  {
+    x = function () {
+      z;
+    };
+  }
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+function component() {
+  const $ = React.useMemoCache();
+  const z = 100;
+  const x = undefined;
+  let x$0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    x$0 = function () {
+      z;
+    };
+
+    $[0] = x$0;
+  } else {
+    x$0 = $[0];
+  }
+
+  return x$0;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/hir/capturing-function-within-block.js

@@ -0,0 +1,10 @@
+function component() {
+  let z = 100;
+  let x;
+  {
+    x = function () {
+      z;
+    };
+  }
+  return x;
+}

compiler/forget/src/__tests__/fixtures/hir/capturing-variable-in-nested-block.expect.md

@@ -0,0 +1,40 @@
+
+## Input
+
+```javascript
+function component() {
+  let z = 100;
+  let x = function () {
+    {
+      z;
+    }
+  };
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+function component() {
+  const $ = React.useMemoCache();
+  const z = 100;
+  let x;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    x = function () {
+      {
+        z;
+      }
+    };
+
+    $[0] = x;
+  } else {
+    x = $[0];
+  }
+
+  return x;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/hir/capturing-variable-in-nested-block.js

@@ -0,0 +1,9 @@
+function component() {
+  let z = 100;
+  let x = function () {
+    {
+      z;
+    }
+  };
+  return x;
+}

compiler/forget/src/__tests__/fixtures/hir/capturing-variable-in-nested-function.expect.md

@@ -0,0 +1,40 @@
+
+## Input
+
+```javascript
+function component() {
+  let z = 100;
+  let x = function () {
+    (function () {
+      z;
+    })();
+  };
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+function component() {
+  const $ = React.useMemoCache();
+  const z = 100;
+  let x;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    x = function () {
+      (function () {
+        z;
+      })();
+    };
+
+    $[0] = x;
+  } else {
+    x = $[0];
+  }
+
+  return x;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/hir/capturing-variable-in-nested-function.js

@@ -0,0 +1,9 @@
+function component() {
+  let z = 100;
+  let x = function () {
+    (function () {
+      z;
+    })();
+  };
+  return x;
+}