From 6f9c9cf9ea037829b4da4bc419a877e3fc8b7d4a Mon Sep 17 00:00:00 2001
From: Sathya Gunasekaran <gsathya.ceg@gmail.com>
Date: Tue, 10 Jan 2023 16:06:55 +0000
Subject: [PATCH] [hir] Make lambdas dependencies be a list of Places

This lets HIR understand and rename dependencies correctly.
---
 compiler/forget/src/HIR/BuildHIR.ts           | 20 +++++--
 compiler/forget/src/HIR/HIR.ts                |  2 +-
 compiler/forget/src/HIR/PrintHIR.ts           |  4 +-
 compiler/forget/src/HIR/visitors.ts           | 10 +++-
 .../hir/prop-capturing-function-1.expect.md   | 53 +++++++++++++++++++
 .../fixtures/hir/prop-capturing-function-1.js |  7 +++
 6 files changed, 87 insertions(+), 9 deletions(-)
 create mode 100644 compiler/forget/src/__tests__/fixtures/hir/prop-capturing-function-1.expect.md
 create mode 100644 compiler/forget/src/__tests__/fixtures/hir/prop-capturing-function-1.js

diff --git a/compiler/forget/src/HIR/BuildHIR.ts b/compiler/forget/src/HIR/BuildHIR.ts
index ac1b308f02..5bee25f325 100644
--- a/compiler/forget/src/HIR/BuildHIR.ts
+++ b/compiler/forget/src/HIR/BuildHIR.ts
@@ -1152,7 +1152,11 @@ function lowerExpression(
       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 dependencies: Array<Place> = gatherCapturedDeps(
+        builder,
+        expr,
+        componentScope
+      );
       const body = expr.get("body").node;
       const params: Array<string> = expr.get("params").map((p) => {
         todoInvariant(p.isIdentifier(), "handle non identifier params");
@@ -1562,10 +1566,11 @@ function capturePureScopes(
 }
 
 function gatherCapturedDeps(
+  builder: HIRBuilder,
   fn: NodePath<t.FunctionExpression>,
   componentScope: Scope
-): Set<t.Identifier> {
-  const captured: Set<t.Identifier> = new Set();
+): Array<Place> {
+  const captured: Set<Place> = new Set();
 
   // Capture all the scopes from the parent of this function up to and including
   // the component scope.
@@ -1587,9 +1592,14 @@ function gatherCapturedDeps(
         return;
       }
 
-      captured.add(binding.identifier);
+      captured.add({
+        kind: "Identifier",
+        identifier: builder.resolveIdentifier(binding.identifier),
+        effect: Effect.Unknown,
+        loc: id.node.loc!,
+      });
     },
   });
 
-  return captured;
+  return [...captured];
 }
diff --git a/compiler/forget/src/HIR/HIR.ts b/compiler/forget/src/HIR/HIR.ts
index bc2d4682f7..075669a31e 100644
--- a/compiler/forget/src/HIR/HIR.ts
+++ b/compiler/forget/src/HIR/HIR.ts
@@ -313,7 +313,7 @@ export type InstructionData =
       kind: "FunctionExpression";
       name: string | null;
       params: Array<string>;
-      dependencies: Set<t.Identifier>;
+      dependencies: Array<Place>;
       body: t.BlockStatement;
     }
 
diff --git a/compiler/forget/src/HIR/PrintHIR.ts b/compiler/forget/src/HIR/PrintHIR.ts
index 19e5ac833e..43115aa036 100644
--- a/compiler/forget/src/HIR/PrintHIR.ts
+++ b/compiler/forget/src/HIR/PrintHIR.ts
@@ -293,7 +293,9 @@ export function printInstructionValue(instrValue: InstructionValue): string {
     case "FunctionExpression": {
       const params = instrValue.params.join(",");
       const body = generate(instrValue.body).code;
-      const deps = [...instrValue.dependencies].map((i) => i.name).join(",");
+      const deps = instrValue.dependencies
+        .map((i) => printIdentifier(i.identifier))
+        .join(",");
       value = `Function ${instrValue.name ?? ""} @deps[${deps}] (${
         params ?? ""
       }){${body}}`;
diff --git a/compiler/forget/src/HIR/visitors.ts b/compiler/forget/src/HIR/visitors.ts
index be172e978f..b09f2011aa 100644
--- a/compiler/forget/src/HIR/visitors.ts
+++ b/compiler/forget/src/HIR/visitors.ts
@@ -87,9 +87,12 @@ export function* eachInstructionValueOperand(
       yield* instrValue.elements;
       break;
     }
+    case "FunctionExpression": {
+      yield* instrValue.dependencies;
+      break;
+    }
     case "OtherStatement":
     case "Primitive":
-    case "FunctionExpression":
     case "JSXText": {
       break;
     }
@@ -174,9 +177,12 @@ export function mapInstructionOperands(
       instrValue.children = instrValue.children.map((e) => fn(e));
       break;
     }
+    case "FunctionExpression": {
+      instrValue.dependencies = instrValue.dependencies.map((d) => fn(d));
+      break;
+    }
     case "OtherStatement":
     case "Primitive":
-    case "FunctionExpression":
     case "JSXText": {
       break;
     }
diff --git a/compiler/forget/src/__tests__/fixtures/hir/prop-capturing-function-1.expect.md b/compiler/forget/src/__tests__/fixtures/hir/prop-capturing-function-1.expect.md
new file mode 100644
index 0000000000..65467a6482
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/hir/prop-capturing-function-1.expect.md
@@ -0,0 +1,53 @@
+
+## Input
+
+```javascript
+function component(a, b) {
+  let z = { a, b };
+  let x = function () {
+    z;
+  };
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+function component(a, b) {
+  const $ = React.useMemoCache();
+  const c_0 = $[0] !== a;
+  const c_1 = $[1] !== b;
+  let z;
+  if (c_0 || c_1) {
+    z = {
+      a: a,
+      b: b,
+    };
+    $[0] = a;
+    $[1] = b;
+    $[2] = z;
+  } else {
+    z = $[2];
+  }
+
+  const c_3 = $[3] !== z;
+  let x;
+
+  if (c_3) {
+    x = function () {
+      z;
+    };
+
+    $[3] = z;
+    $[4] = x;
+  } else {
+    x = $[4];
+  }
+
+  return x;
+}
+
+```
+      
\ No newline at end of file
diff --git a/compiler/forget/src/__tests__/fixtures/hir/prop-capturing-function-1.js b/compiler/forget/src/__tests__/fixtures/hir/prop-capturing-function-1.js
new file mode 100644
index 0000000000..0817daa5d1
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/hir/prop-capturing-function-1.js
@@ -0,0 +1,7 @@
+function component(a, b) {
+  let z = { a, b };
+  let x = function () {
+    z;
+  };
+  return x;
+}