[codegen] location refinements

- Add a `loc` to the `while` terminal node. 

- Move location data for assignments 1 level higher as that seems to work better 
in the generated code (not tested with actual debugger yet though, we'll 
probably want to look at this more closely.

compiler/forget/src/HIR/BuildHIR.ts

@@ -465,9 +465,12 @@ function lowerStatement(
        * The code leading up to the loop must jump to the conditional block,
        * to evaluate whether to enter the loop or bypass to the continuation.
        */
+      const loc = stmt.node.loc;
+      invariant(loc, "while statement must have a location");
       builder.terminateWithContinuation(
         {
           kind: "while",
+          loc,
           test: conditionalBlock.id,
           loop: loopBlock,
           fallthrough: continuationBlock.id,

compiler/forget/src/HIR/Codegen.ts

@@ -25,10 +25,16 @@ 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> => {
+): (
+  loc: SourceLocation | null | undefined,
+  ...args: Parameters<T>
+) => ReturnType<T> {
+  return (
+    loc: SourceLocation | null | undefined,
+    ...args: Parameters<T>
+  ): ReturnType<T> => {
     const node = fn(...args);
-    if (loc != GeneratedSource) {
+    if (loc != null && loc != GeneratedSource) {
       node.loc = loc;
     }
     // @ts-ignore
@@ -36,9 +42,11 @@ function withLoc<TNode extends t.Node, T extends (...args: any[]) => TNode>(
   };
 }
 
+const createExpressionStatement = withLoc(t.expressionStatement);
 const createFunctionDeclaration = withLoc(t.functionDeclaration);
+const createLabelledStatement = withLoc(t.labeledStatement);
 const createVariableDeclaration = withLoc(t.variableDeclaration);
-const createAssignmentExpression = withLoc(t.assignmentExpression);
+const createWhileStatement = withLoc(t.whileStatement);
 
 /**
  * Converts HIR into Babel nodes, which can then be printed into source text.
@@ -116,7 +124,7 @@ class CodegenVisitor
     } else {
       switch (instr.lvalue.kind) {
         case InstructionKind.Const: {
-          return t.variableDeclaration("const", [
+          return createVariableDeclaration(instr.loc, "const", [
             t.variableDeclarator(codegenLVal(instr.lvalue), value),
           ]);
         }
@@ -126,13 +134,9 @@ class CodegenVisitor
           ]);
         }
         case InstructionKind.Reassign: {
-          return t.expressionStatement(
-            createAssignmentExpression(
-              instr.loc,
-              "=",
-              codegenLVal(instr.lvalue),
-              value
-            )
+          return createExpressionStatement(
+            instr.loc,
+            t.assignmentExpression("=", codegenLVal(instr.lvalue), value)
           );
         }
         default: {
@@ -181,7 +185,7 @@ class CodegenVisitor
         return t.switchStatement(terminal.test, terminal.cases);
       }
       case "while": {
-        return t.whileStatement(terminal.test, terminal.loop);
+        return createWhileStatement(terminal.loc, terminal.test, terminal.loop);
       }
       case "return": {
         if (terminal.value !== null) {
@@ -213,7 +217,7 @@ class CodegenVisitor
       return;
     }
     if (label !== undefined) {
-      block.push(t.labeledStatement(t.identifier(label), item));
+      block.push(createLabelledStatement(item.loc, t.identifier(label), item));
     } else {
       block.push(item);
     }

compiler/forget/src/HIR/HIR.ts

@@ -160,6 +160,7 @@ export type SwitchTerminal = {
 
 export type WhileTerminal = {
   kind: "while";
+  loc: SourceLocation;
   test: BlockId;
   loop: BlockId;
   fallthrough: BlockId;

compiler/forget/src/HIR/HIRTreeVisitor.ts

@@ -16,6 +16,7 @@ import {
   Instruction,
   InstructionValue,
   Place,
+  SourceLocation,
 } from "./HIR";
 
 /**
@@ -270,6 +271,7 @@ class Driver<TBlock, TValue, TItem, TCase> {
             blockValue,
             this.visitor.visitTerminal({
               kind: "while",
+              loc: terminal.loc,
               test: testValue,
               loop: loopBody,
             }),
@@ -281,6 +283,7 @@ class Driver<TBlock, TValue, TItem, TCase> {
             blockValue,
             this.visitor.visitTerminal({
               kind: "while",
+              loc: terminal.loc,
               test: testValue,
               loop: loopBody,
             })
@@ -679,6 +682,6 @@ export type BlockTerminal<TBlock, TValue, TItem, TCase> =
       alternate: TItem | null;
     }
   | { kind: "switch"; test: TValue; cases: Array<TCase> }
-  | { kind: "while"; test: TValue; loop: TItem }
+  | { kind: "while"; loc: SourceLocation; test: TValue; loop: TItem }
   | { kind: "break"; label: string | null }
   | { kind: "continue"; label: string | null };

compiler/forget/src/HIR/visitors.ts

@@ -201,6 +201,7 @@ export function mapTerminalSuccessors(
       const fallthrough = fn(terminal.fallthrough, true);
       return {
         kind: "while",
+        loc: terminal.loc,
         test,
         loop,
         fallthrough,