From 1228bc8ce5eadd15a4d25ba373afcb4eef8543be Mon Sep 17 00:00:00 2001
From: Jan Kassens <jkassens@meta.com>
Date: Mon, 21 Nov 2022 15:14:20 -0500
Subject: [PATCH] [codegen] assign some locations in generated code (#802)

---
 .../components/Editor/HIRTabContent.tsx       |  8 ++---
 compiler/forget/src/HIR/Codegen.ts            | 31 +++++++++++++++++--
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/compiler/forget/packages/playground/components/Editor/HIRTabContent.tsx b/compiler/forget/packages/playground/components/Editor/HIRTabContent.tsx
index fcf8350b11..c0f60d7d4f 100644
--- a/compiler/forget/packages/playground/components/Editor/HIRTabContent.tsx
+++ b/compiler/forget/packages/playground/components/Editor/HIRTabContent.tsx
@@ -38,7 +38,7 @@ export default function HIRTabContent({ source }: { source: string }) {
   });
 
   return (
-    <div className="h-full flex flex-col">
+    <div className="flex flex-col h-full">
       <div className="flex-1">
         {astFunctions.map((func, index): React.ReactNode => {
           let body;
@@ -61,8 +61,6 @@ export default function HIRTabContent({ source }: { source: string }) {
 
             if (flags.codegen) {
               const ast = codegen(ir);
-              // TODO: codegen should set this correctly directly
-              ast.loc = func.node.loc;
               const generated = generate(
                 ast,
                 {
@@ -100,8 +98,8 @@ export default function HIRTabContent({ source }: { source: string }) {
           }
           const name = func.node.id?.name ?? "anonymous";
           return (
-            <div key={index} className="border-b py-4 px-2">
-              <h3 className="font-medium mb-2">Function: {name}</h3>
+            <div key={index} className="px-2 py-4 border-b">
+              <h3 className="mb-2 font-medium">Function: {name}</h3>
               {body}
             </div>
           );
diff --git a/compiler/forget/src/HIR/Codegen.ts b/compiler/forget/src/HIR/Codegen.ts
index 716c299522..0fc01752b0 100644
--- a/compiler/forget/src/HIR/Codegen.ts
+++ b/compiler/forget/src/HIR/Codegen.ts
@@ -9,6 +9,7 @@ import * as t from "@babel/types";
 import { assertExhaustive } from "../Common/utils";
 import { invariant } from "../CompilerError";
 import {
+  GeneratedSource,
   HIRFunction,
   Identifier,
   IdentifierId,
@@ -17,10 +18,28 @@ import {
   InstructionValue,
   LValue,
   Place,
+  SourceLocation,
 } from "./HIR";
 import { BlockTerminal, Visitor, visitTree } from "./HIRTreeVisitor";
 import { todoInvariant } from "./todo";
 
+function withLoc<TNode extends t.Node, T extends (...args: any[]) => TNode>(
+  fn: T
+): (loc: SourceLocation, ...args: Parameters<T>) => ReturnType<T> {
+  return (loc: SourceLocation, ...args: Parameters<T>): ReturnType<T> => {
+    const node = fn(...args);
+    if (loc != GeneratedSource) {
+      node.loc = loc;
+    }
+    // @ts-ignore
+    return node;
+  };
+}
+
+const createFunctionDeclaration = withLoc(t.functionDeclaration);
+const createVariableDeclaration = withLoc(t.variableDeclaration);
+const createAssignmentExpression = withLoc(t.assignmentExpression);
+
 /**
  * Converts HIR into Babel nodes, which can then be printed into source text.
  * Note that converting source to HIR and back is not guaranteed to produce
@@ -46,7 +65,8 @@ export default function codegen(fn: HIRFunction): t.Function {
   const body = visitTree(fn, visitor);
   invariant(t.isBlockStatement(body), "Expected a block statement");
   const params = fn.params.map((param) => convertIdentifier(param.identifier));
-  return t.functionDeclaration(
+  return createFunctionDeclaration(
+    fn.loc,
     fn.id !== null ? convertIdentifier(fn.id) : null,
     params,
     body,
@@ -101,13 +121,18 @@ class CodegenVisitor
           ]);
         }
         case InstructionKind.Let: {
-          return t.variableDeclaration("let", [
+          return createVariableDeclaration(instr.loc, "let", [
             t.variableDeclarator(codegenLVal(instr.lvalue), value),
           ]);
         }
         case InstructionKind.Reassign: {
           return t.expressionStatement(
-            t.assignmentExpression("=", codegenLVal(instr.lvalue), value)
+            createAssignmentExpression(
+              instr.loc,
+              "=",
+              codegenLVal(instr.lvalue),
+              value
+            )
           );
         }
         default: {