fix break

src/ast.ts

@@ -127,19 +127,17 @@ export type ExprIf = {
   else?: Expr;
 };
 
+export type LoopId = number;
+
 export type ExprLoop = {
   kind: "loop";
   body: Expr;
+  loopId: LoopId;
 };
 
 export type ExprBreak = {
   kind: "break";
-  /**
-   * The break target block.
-   * May be any control flow block, labelled from inside out.
-   * TODO: This is not a good solution at all and pretty broken.
-   */
-  target?: number;
+  target?: LoopId;
 };
 
 export type ExprStructLiteral = {

src/index.ts

@@ -12,6 +12,7 @@ import { exec } from "child_process";
 const input = `
 function main() = (
  let i = 0;
+ loop break;
  loop (
   if i > 10 then break;
 

src/lower.ts

@@ -5,12 +5,13 @@ import {
   FunctionDef,
   ImportDef,
   Item,
+  LoopId,
   Resolution,
   Ty,
   TyFn,
   varUnreachable,
 } from "./ast";
-import { encodeUtf8 } from "./utils";
+import { encodeUtf8, unwrap } from "./utils";
 import * as wasm from "./wasm/defs";
 
 const USIZE: wasm.ValType = "i32";
@@ -204,6 +205,8 @@ type FuncContext = {
   func: FunctionDef;
   wasm: wasm.Func;
   varLocations: VarLocation[];
+  loopDepths: Map<LoopId, number>;
+  currentBlockDepth: number;
 };
 
 type FnAbi = { params: ArgRetAbi[]; ret: ArgRetAbi };
@@ -230,6 +233,8 @@ function lowerFunc(cx: Context, item: Item, func: FunctionDef) {
     func,
     wasm: wasmFunc,
     varLocations: paramLocations,
+    loopDepths: new Map(),
+    currentBlockDepth: 0,
   };
 
   lowerExpr(fcx, wasmFunc.body, fcx.func.body);
@@ -545,6 +550,7 @@ function lowerExpr(fcx: FuncContext, instrs: wasm.Instr[], expr: Expr) {
     case "if": {
       lowerExpr(fcx, instrs, expr.cond!);
 
+      fcx.currentBlockDepth++;
       const thenInstrs: wasm.Instr[] = [];
       lowerExpr(fcx, thenInstrs, expr.then);
 
@@ -553,6 +559,7 @@ function lowerExpr(fcx: FuncContext, instrs: wasm.Instr[], expr: Expr) {
       if (expr.else) {
         lowerExpr(fcx, elseInstrs, expr.else);
       }
+      fcx.currentBlockDepth--;
 
       instrs.push({
         kind: "if",
@@ -566,12 +573,15 @@ function lowerExpr(fcx: FuncContext, instrs: wasm.Instr[], expr: Expr) {
     case "loop": {
       const outerBlockInstrs: wasm.Instr[] = [];
 
+      fcx.loopDepths.set(expr.loopId, fcx.currentBlockDepth);
+      fcx.currentBlockDepth += 2;
       const bodyInstrs: wasm.Instr[] = [];
       lowerExpr(fcx, bodyInstrs, expr.body);
       bodyInstrs.push({
         kind: "br",
         label: /*innermost control structure, the loop*/ 0,
       });
+      fcx.currentBlockDepth -= 2;
 
       outerBlockInstrs.push({
         kind: "loop",
@@ -584,13 +594,18 @@ function lowerExpr(fcx: FuncContext, instrs: wasm.Instr[], expr: Expr) {
         instrs: outerBlockInstrs,
         type: blockTypeForBody(fcx.cx, expr.ty!),
       });
+      fcx.loopDepths.delete(expr.loopId);
 
       break;
     }
     case "break": {
+      const loopDepth = unwrap(fcx.loopDepths.get(expr.target!));
+
+      console.log(loopDepth, fcx.currentBlockDepth);
+
       instrs.push({
         kind: "br",
-        label: expr.target! + /* the block outside the loop */ 1,
+        label: fcx.currentBlockDepth - loopDepth - 1,
       });
       break;
     }
@@ -612,6 +627,7 @@ function lowerExprBlockBody(
   fcx: FuncContext,
   expr: ExprBlock & Expr
 ): wasm.Instr[] {
+  fcx.currentBlockDepth++;
   const innerInstrs: wasm.Instr[] = [];
 
   const headExprs = expr.exprs.slice(0, -1);
@@ -625,6 +641,8 @@ function lowerExprBlockBody(
 
   lowerExpr(fcx, innerInstrs, tailExpr);
 
+  fcx.currentBlockDepth--;
+
   return innerInstrs;
 }
 

src/parser.ts

@@ -6,6 +6,7 @@ import {
   COMPARISON_KINDS,
   DEFAULT_FOLDER,
   Expr,
+  ExprLoop,
   ExprStructLiteral,
   FieldDef,
   Folder,
@@ -24,13 +25,8 @@ import {
   superFoldExpr,
 } from "./ast";
 import { CompilerError, Span, spanMerge } from "./error";
-import {
-  BaseToken,
-  Token,
-  TokenIdent,
-  TokenLit,
-  TokenLitString,
-} from "./lexer";
+import { BaseToken, Token, TokenIdent, TokenLitString } from "./lexer";
+import { Ids } from "./utils";
 
 type Parser<T> = (t: Token[]) => [Token[], T];
 
@@ -43,9 +39,7 @@ export function parse(t: Token[]): Ast {
     items.push(item);
   }
 
-  const withIds = items.map((item, i) => ({ ...item, id: i }));
-
-  const ast = { items: withIds };
+  const ast = assignIds({ items: items });
 
   validateAst(ast);
 
@@ -406,7 +400,7 @@ function parseExprAtom(startT: Token[]): [Token[], Expr] {
   if (tok.kind === "loop") {
     let body;
     [t, body] = parseExpr(t);
-    return [t, { kind: "loop", body, span: tok.span }];
+    return [t, { kind: "loop", body, span: tok.span, loopId: 0 }];
   }
 
   if (tok.kind === "break") {
@@ -614,3 +608,24 @@ function validateAst(ast: Ast) {
 
   foldAst(ast, validator);
 }
+
+function assignIds(ast: Ast): Ast {
+  let loopId = new Ids();
+
+  const astItems = { items: ast.items.map((item, i) => ({ ...item, id: i })) };
+
+  const assigner: Folder = {
+    ...DEFAULT_FOLDER,
+    expr(expr: Expr): Expr {
+      if (expr.kind === "loop") {
+        return {
+          ...(superFoldExpr(expr, this) as ExprLoop & Expr),
+          loopId: loopId.next(),
+        };
+      }
+      return superFoldExpr(expr, this);
+    },
+  };
+
+  return foldAst(astItems, assigner);
+}

src/typeck.ts

@@ -13,6 +13,7 @@ import {
   Identifier,
   ItemId,
   LOGICAL_KINDS,
+  LoopId,
   Resolution,
   Ty,
   TY_BOOL,
@@ -483,8 +484,7 @@ export function checkBody(
   typeOfItem: (index: number) => Ty
 ): Expr {
   const localTys = [...fnTy.params];
-  const loopState: { hasBreak: boolean; nestingDepth: number }[] = [];
-  let currentNestingDepth = 0;
+  const loopState: { hasBreak: boolean; loopId: LoopId }[] = [];
 
   const infcx = new InferContext();
 
@@ -586,7 +586,6 @@ export function checkBody(
           };
         }
         case "block": {
-          currentNestingDepth++;
           const prevLocalTysLen = localTys.length;
 
           const exprs = expr.exprs.map((expr) => this.expr(expr));
@@ -595,8 +594,6 @@ export function checkBody(
 
           localTys.length = prevLocalTysLen;
 
-          currentNestingDepth--;
-
           return {
             ...expr,
             exprs,
@@ -680,10 +677,8 @@ export function checkBody(
         }
         case "if": {
           const cond = this.expr(expr.cond);
-          currentNestingDepth++;
           const then = this.expr(expr.then);
           const elsePart = expr.else && this.expr(expr.else);
-          currentNestingDepth--;
 
           infcx.assign(TY_BOOL, cond.ty!, cond.span);
 
@@ -699,10 +694,9 @@ export function checkBody(
           return { ...expr, cond, then, else: elsePart, ty };
         }
         case "loop": {
-          currentNestingDepth++;
           loopState.push({
             hasBreak: false,
-            nestingDepth: currentNestingDepth,
+            loopId: expr.loopId,
           });
 
           const body = this.expr(expr.body);
@@ -711,8 +705,6 @@ export function checkBody(
           const hadBreak = loopState.pop();
           const ty = hadBreak ? TY_UNIT : TY_NEVER;
 
-          currentNestingDepth--;
-
           return {
             ...expr,
             body,
@@ -723,11 +715,9 @@ export function checkBody(
           if (loopState.length === 0) {
             throw new CompilerError("break outside loop", expr.span);
           }
-          const loopDepth = loopState[loopState.length - 1].nestingDepth;
+          const target = loopState[loopState.length - 1].loopId;
           loopState[loopState.length - 1].hasBreak = true;
 
-          const target = currentNestingDepth - loopDepth;
-
           return {
             ...expr,
             ty: TY_NEVER,

src/utils.ts

@@ -1,3 +1,18 @@
 export function encodeUtf8(s: string): Uint8Array {
   return new TextEncoder().encode(s);
 }
+
+export class Ids {
+  nextId: number = 0;
+
+  public next(): number {
+    return this.nextId++;
+  }
+}
+
+export function unwrap<T>(value: T | undefined): T {
+  if (value === undefined) {
+    throw new Error("tried to unwrap undefined value");
+  }
+  return value;
+}