From e2e77b4b06e51c7f7d3ea187defaf1ad08e513c1 Mon Sep 17 00:00:00 2001
From: Mathias Magnusson <mathias@magnusson.space>
Date: Mon, 2 Jun 2025 21:55:11 +0200
Subject: 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)
---
 src/codegen.zig | 26 ++++++++++++++------------
 1 file changed, 14 insertions(+), 12 deletions(-)

(limited to 'src/codegen.zig')

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| {
-- 
cgit v1.2.3