From 3927281fe4da48f0d9eb78019f832fd404a7c089 Mon Sep 17 00:00:00 2001
From: Sathya Gunasekaran <gsathya.ceg@gmail.com>
Date: Mon, 23 Jan 2023 13:20:33 -0500
Subject: [PATCH] =?UTF-8?q?[=CE=BB]=20Re=20run=20SSA=20on=20function=20exp?=
 =?UTF-8?q?r?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Functions can capture variables declared after the definition of the function.
This re runs SSA to map the captured identifiers to the new SSA identifiers if
available.
---
 compiler/forget/src/SSA/EnterSSA.ts           |  8 ++++
 .../mutate-captured-arg-separately.expect.md  | 38 +++++++++++++++++++
 .../hir/mutate-captured-arg-separately.js     |  9 +++++
 3 files changed, 55 insertions(+)
 create mode 100644 compiler/forget/src/__tests__/fixtures/hir/mutate-captured-arg-separately.expect.md
 create mode 100644 compiler/forget/src/__tests__/fixtures/hir/mutate-captured-arg-separately.js

diff --git a/compiler/forget/src/SSA/EnterSSA.ts b/compiler/forget/src/SSA/EnterSSA.ts
index e35f3664a4..0f6f9ec5c1 100644
--- a/compiler/forget/src/SSA/EnterSSA.ts
+++ b/compiler/forget/src/SSA/EnterSSA.ts
@@ -226,4 +226,12 @@ export default function enterSSA(func: HIRFunction) {
       }
     }
   }
+
+  for (const [_, block] of func.body.blocks) {
+    for (const instr of block.instructions) {
+      if (instr.value.kind === "FunctionExpression") {
+        mapInstructionOperands(instr, (place) => builder.getPlace(place));
+      }
+    }
+  }
 }
diff --git a/compiler/forget/src/__tests__/fixtures/hir/mutate-captured-arg-separately.expect.md b/compiler/forget/src/__tests__/fixtures/hir/mutate-captured-arg-separately.expect.md
new file mode 100644
index 0000000000..5745524c79
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/hir/mutate-captured-arg-separately.expect.md
@@ -0,0 +1,38 @@
+
+## Input
+
+```javascript
+function component(a) {
+  let y = function () {
+    m(x);
+  };
+
+  let x = { a };
+  m(x);
+  return y;
+}
+
+```
+
+## Code
+
+```javascript
+function component(a) {
+  const $ = React.useMemoCache();
+  let y;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    y = function () {
+      m(x);
+    };
+    $[0] = y;
+  } else {
+    y = $[0];
+  }
+
+  const x = { a: a };
+  m(x);
+  return y;
+}
+
+```
+      
\ No newline at end of file
diff --git a/compiler/forget/src/__tests__/fixtures/hir/mutate-captured-arg-separately.js b/compiler/forget/src/__tests__/fixtures/hir/mutate-captured-arg-separately.js
new file mode 100644
index 0000000000..5e083ec453
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/hir/mutate-captured-arg-separately.js
@@ -0,0 +1,9 @@
+function component(a) {
+  let y = function () {
+    m(x);
+  };
+
+  let x = { a };
+  m(x);
+  return y;
+}