From ee3a88f20f773c91c538420a420ae2435277976f Mon Sep 17 00:00:00 2001
From: Sathya Gunasekaran <gsathya.ceg@gmail.com>
Date: Mon, 23 Jan 2023 17:39:37 -0500
Subject: [PATCH] =?UTF-8?q?[=CE=BB]=20Parse=20and=20codegen=20arrow=20func?=
 =?UTF-8?q?tion=20exprs?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 compiler/forget/src/HIR/BuildHIR.ts           | 15 ++++---
 compiler/forget/src/HIR/HIR.ts                |  2 +-
 compiler/forget/src/HIR/PrintHIR.ts           |  2 +-
 .../ReactiveScopes/CodegenReactiveFunction.ts |  2 +-
 .../hir/capturing-arrow-function-1.expect.md  | 44 +++++++++++++++++++
 .../hir/capturing-arrow-function-1.js         |  7 +++
 6 files changed, 64 insertions(+), 8 deletions(-)
 create mode 100644 compiler/forget/src/__tests__/fixtures/hir/capturing-arrow-function-1.expect.md
 create mode 100644 compiler/forget/src/__tests__/fixtures/hir/capturing-arrow-function-1.js

diff --git a/compiler/forget/src/HIR/BuildHIR.ts b/compiler/forget/src/HIR/BuildHIR.ts
index 5e8ed1e37c..0b0346e30f 100644
--- a/compiler/forget/src/HIR/BuildHIR.ts
+++ b/compiler/forget/src/HIR/BuildHIR.ts
@@ -1334,16 +1334,21 @@ function lowerExpression(
         loc: exprLoc,
       };
     }
+    case "ArrowFunctionExpression":
     case "FunctionExpression": {
-      const expr = exprPath as NodePath<t.FunctionExpression>;
-      const name: string | null = expr.get("id")?.node?.name ?? null;
+      const expr = exprPath as NodePath<
+        t.FunctionExpression | t.ArrowFunctionExpression
+      >;
+      let name: string | null = null;
+      if (expr.isFunctionExpression()) {
+        name = expr.get("id")?.node?.name ?? null;
+      }
       const componentScope: Scope = expr.scope.parent.getFunctionParent()!;
       const dependencies: Array<Place> = gatherCapturedDeps(
         builder,
         expr,
         componentScope
       );
-      const body = expr.get("body").node;
       const lowering = lower(expr);
       let loweredFunc: HIRFunction;
       if (lowering.isErr()) {
@@ -1375,11 +1380,11 @@ function lowerExpression(
         : {
             kind: "FunctionExpression",
             name,
-            body,
             params,
             loweredFunc,
             dependencies,
             mutatedDeps: [],
+            expr: expr.node,
             loc: exprLoc,
           };
     }
@@ -1901,7 +1906,7 @@ function captureScopes({ from, to }: { from: Scope; to: Scope }): Set<Scope> {
 
 function gatherCapturedDeps(
   builder: HIRBuilder,
-  fn: NodePath<t.FunctionExpression>,
+  fn: NodePath<t.FunctionExpression | t.ArrowFunctionExpression>,
   componentScope: Scope
 ): Array<Place> {
   const captured: Set<Place> = new Set();
diff --git a/compiler/forget/src/HIR/HIR.ts b/compiler/forget/src/HIR/HIR.ts
index 47078c75b9..574cbdc6df 100644
--- a/compiler/forget/src/HIR/HIR.ts
+++ b/compiler/forget/src/HIR/HIR.ts
@@ -407,7 +407,7 @@ export type InstructionData =
       // source of truth.
       mutatedDeps: Array<Place>;
       loweredFunc: HIRFunction;
-      body: t.BlockStatement;
+      expr: t.ArrowFunctionExpression | t.FunctionExpression;
     }
 
   /**
diff --git a/compiler/forget/src/HIR/PrintHIR.ts b/compiler/forget/src/HIR/PrintHIR.ts
index 505895c7f9..b81dad85e6 100644
--- a/compiler/forget/src/HIR/PrintHIR.ts
+++ b/compiler/forget/src/HIR/PrintHIR.ts
@@ -312,7 +312,7 @@ export function printInstructionValue(instrValue: InstructionValue): string {
     }
     case "FunctionExpression": {
       const params = instrValue.params.join(",");
-      const body = generate(instrValue.body).code;
+      const body = generate(instrValue.expr).code;
       const deps = instrValue.dependencies
         .map((i) => printIdentifier(i.identifier))
         .join(",");
diff --git a/compiler/forget/src/ReactiveScopes/CodegenReactiveFunction.ts b/compiler/forget/src/ReactiveScopes/CodegenReactiveFunction.ts
index 1ab07d09c5..17968a77a5 100644
--- a/compiler/forget/src/ReactiveScopes/CodegenReactiveFunction.ts
+++ b/compiler/forget/src/ReactiveScopes/CodegenReactiveFunction.ts
@@ -662,7 +662,7 @@ function codegenInstructionValue(
       const id =
         instrValue.name !== null ? t.identifier(instrValue.name) : null;
       const params = instrValue.params.map((p) => t.identifier(p));
-      value = t.functionExpression(id, params, instrValue.body);
+      value = instrValue.expr;
       break;
     }
     default: {
diff --git a/compiler/forget/src/__tests__/fixtures/hir/capturing-arrow-function-1.expect.md b/compiler/forget/src/__tests__/fixtures/hir/capturing-arrow-function-1.expect.md
new file mode 100644
index 0000000000..8271d8fdae
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/hir/capturing-arrow-function-1.expect.md
@@ -0,0 +1,44 @@
+
+## Input
+
+```javascript
+function component(a) {
+  let z = { a };
+  let x = () => {
+    z;
+  };
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+function component(a) {
+  const $ = React.useMemoCache();
+  const c_0 = $[0] !== a;
+  let z;
+  if (c_0) {
+    z = { a: a };
+    $[0] = a;
+    $[1] = z;
+  } else {
+    z = $[1];
+  }
+  const c_2 = $[2] !== z;
+  let x;
+  if (c_2) {
+    x = () => {
+      z;
+    };
+    $[2] = z;
+    $[3] = x;
+  } else {
+    x = $[3];
+  }
+  return x;
+}
+
+```
+      
\ No newline at end of file
diff --git a/compiler/forget/src/__tests__/fixtures/hir/capturing-arrow-function-1.js b/compiler/forget/src/__tests__/fixtures/hir/capturing-arrow-function-1.js
new file mode 100644
index 0000000000..c7efd24dfb
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/hir/capturing-arrow-function-1.js
@@ -0,0 +1,7 @@
+function component(a) {
+  let z = { a };
+  let x = () => {
+    z;
+  };
+  return x;
+}