remove the need for explicit discard instructions

by also considering an instruction's destination register(s) to be uses
and when cleaning up registers, also cleaning up those from previous
instructions (specifically this is the output of the last instruction if
it is not used anywhere)

1 files changed, 14 insertions, 12 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index d2b5d2a..f186304 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -560,11 +560,16 @@ const Context = struct {
         try self.instructions.append(inst);
     }
 
-    fn maybeFreeSources(self: *Context) !void {
-        const index = self.current_instruction_index.?;
-        if (self.block.?.vreg_last_use.get(index)) |last_used_here| {
-            for (last_used_here.items) |vreg| {
-                self.register_allocator.free(vreg);
+    /// Frees all virtual registers who's last use is the current `compile.Instr` or earlier. This
+    /// must be called after the current compile.Instr's sources have been retrieved, since this
+    /// will deallocate them, and after allocating auxiliary registers since otherwise they may
+    /// collide with the sources. Should be called before allocating results to allow for more
+    /// register re-use.
+    fn freeUnusedVRegs(self: *Context) !void {
+        var it = self.register_allocator.allocated.keyIterator();
+        while (it.next()) |vreg| {
+            if (self.block.?.vreg_last_use.get(vreg.*).? <= self.current_instruction_index.?) {
+                self.register_allocator.free(vreg.*);
             }
         }
     }
@@ -595,6 +600,7 @@ const Context = struct {
     }
 
     fn genConstant(self: *Context, constant: compile.Instr.Constant) !void {
+        try self.freeUnusedVRegs();
         const reg = try self.register_allocator.allocate(constant.dest);
         try self.genConstantInner(reg, constant.value);
     }
@@ -602,7 +608,7 @@ const Context = struct {
     fn genBinOp(self: *Context, bin_op: compile.Instr.BinOp) !void {
         const lhs = self.register_allocator.get(bin_op.lhs);
         const rhs = self.register_allocator.get(bin_op.rhs);
-        try self.maybeFreeSources();
+        try self.freeUnusedVRegs();
         const reg = try self.register_allocator.allocate(bin_op.dest);
         switch (bin_op.op) {
             .add => try self.emit(.add(reg, lhs, rhs)),
@@ -621,7 +627,7 @@ const Context = struct {
                 defer self.register_allocator.freeAux(digit);
                 const count = try self.register_allocator.allocate(call.dest);
 
-                try self.maybeFreeSources();
+                try self.freeUnusedVRegs();
 
                 try self.emit(.addi(digit, .zero, '\n'));
                 try self.emit(.addi(.sp, .sp, -1));
@@ -648,7 +654,7 @@ const Context = struct {
                 try self.emit(.addi(count, count, -1));
             },
             .read_int => {
-                try self.maybeFreeSources();
+                try self.freeUnusedVRegs();
 
                 const result = try self.register_allocator.allocate(call.dest);
                 const ptr = try self.register_allocator.allocateAux();
@@ -689,10 +695,6 @@ const Context = struct {
         }
     }
 
-    fn genDiscard(self: *Context, _: compile.Instr.Discard) !void {
-        try self.maybeFreeSources();
-    }
-
     fn codegenInstr(self: *Context, instr: compile.Instr) !void {
         switch (instr.type) {
             inline else => |ty| {