bugfix: don't add multiple exits to basic blocks

1 files changed, 42 insertions, 19 deletions

diff --git a/src/compile.zig b/src/compile.zig
index 80279fb..d1cac8c 100644
--- a/src/compile.zig
+++ b/src/compile.zig
@@ -98,14 +98,18 @@ pub const BasicBlock = struct {
 
     fn finalize(self: *BasicBlock, allocator: Allocator) !void {
         std.debug.assert(self.instrs.items.len > 0);
-        for (self.instrs.items[0 .. self.instrs.items.len - 1]) |instr| {
-            std.debug.assert(switch (instr.type) {
-                inline else => |ty| !@hasDecl(@TypeOf(ty), "ends_block"),
-            });
+        for (self.instrs.items, 0..) |instr, i| {
+            const is_last = i == self.instrs.items.len - 1;
+            if (instr.ends_block() != is_last) {
+                std.log.debug("Instruction at {s}last position must {s}end block. Found instruction {}. Block is: {}", .{
+                    if (is_last) "" else "not ",
+                    if (is_last) "" else "not ",
+                    instr,
+                    self,
+                });
+                unreachable;
+            }
         }
-        std.debug.assert(switch (self.instrs.getLast().type) {
-            inline else => |ty| @hasDecl(@TypeOf(ty), "ends_block"),
-        });
 
         self.vreg_last_use = .empty;
         var set_at: std.AutoHashMapUnmanaged(VReg, usize) = .empty;
@@ -155,6 +159,12 @@ pub const Instr = struct {
         ret: Return,
     };
 
+    fn ends_block(self: Instr) bool {
+        return switch (self.type) {
+            inline else => |t| @hasDecl(@TypeOf(t), "ends_block"),
+        };
+    }
+
     pub const Constant = struct {
         dest: VReg,
         value: u64,
@@ -389,15 +399,17 @@ fn compileProcedure(
         break :blk lvar;
     } else null;
     try pctx.compileBlock(proc.body);
-    const proc_res = pctx.vreg_ctr.get();
-    try pctx.addInstr(.{
-        .loc = .{ .start = 0, .end = 0 },
-        .type = .{ .constant = .{ .dest = proc_res, .value = 0 } },
-    });
-    try pctx.addInstr(.{
-        .loc = .{ .start = 0, .end = 0 },
-        .type = .{ .ret = .{ .val = proc_res } },
-    });
+    if (!pctx.currentBlockEnded()) {
+        const proc_res = pctx.vreg_ctr.get();
+        try pctx.addInstr(.{
+            .loc = .{ .start = 0, .end = 0 },
+            .type = .{ .constant = .{ .dest = proc_res, .value = 0 } },
+        });
+        try pctx.addInstr(.{
+            .loc = .{ .start = 0, .end = 0 },
+            .type = .{ .ret = .{ .val = proc_res } },
+        });
+    }
     var blocks = try ctx.allocator.alloc(BasicBlock, pctx.blocks.count());
     var it = pctx.blocks.iterator();
     var i: usize = 0;
@@ -496,7 +508,7 @@ const ProcedureContext = struct {
                 });
                 self.current_block = do;
                 try self.compileBlock(@"while".do);
-                try self.addInstr(.{
+                try self.addInstrUnlessEnded(.{
                     .loc = stmt.loc,
                     .type = .{ .jump = .{ .to = cond_block } },
                 });
@@ -579,7 +591,7 @@ const ProcedureContext = struct {
 
                 const t = try self.switchToNewBlock();
                 try self.compileBlock(@"if".then);
-                try self.addInstr(.{
+                try self.addInstrUnlessEnded(.{
                     .loc = expr.loc,
                     .type = .{ .jump = .{ .to = after } },
                 });
@@ -587,7 +599,7 @@ const ProcedureContext = struct {
                 const f = if (@"if".@"else") |@"else"| blk: {
                     const f = try self.switchToNewBlock();
                     try self.compileBlock(@"else");
-                    try self.addInstr(.{
+                    try self.addInstrUnlessEnded(.{
                         .loc = expr.loc,
                         .type = .{ .jump = .{ .to = after } },
                     });
@@ -632,4 +644,15 @@ const ProcedureContext = struct {
         try self.blocks.getPtr(self.current_block).?
             .instrs.append(self.ctx.allocator, instr);
     }
+
+    fn addInstrUnlessEnded(self: *ProcedureContext, instr: Instr) !void {
+        if (!self.currentBlockEnded()) try self.addInstr(instr);
+    }
+
+    fn currentBlockEnded(self: *ProcedureContext) bool {
+        return if (self.blocks.get(self.current_block).?.instrs.getLastOrNull()) |instr|
+            instr.ends_block()
+        else
+            false;
+    }
 };