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)

3 files changed, 22 insertions, 34 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| {
diff --git a/src/compile.zig b/src/compile.zig
index c837495..ac7768e 100644
--- a/src/compile.zig
+++ b/src/compile.zig
@@ -15,7 +15,6 @@ pub const Instr = struct {
         constant: Constant,
         bin_op: BinOp,
         proc_call: ProcCall,
-        discard: Discard,
     };
 
     pub const Constant = struct {
@@ -58,14 +57,6 @@ pub const Instr = struct {
         }
     };
 
-    pub const Discard = struct {
-        vreg: VReg,
-
-        pub fn sources(self: Discard) Sources {
-            return Sources.fromSlice(&.{self.vreg}) catch unreachable;
-        }
-    };
-
     pub fn sources(self: Instr) Sources {
         return switch (self.type) {
             inline else => |instr| instr.sources(),
@@ -75,7 +66,6 @@ pub const Instr = struct {
     pub fn dest(self: *const Instr) ?VReg {
         return switch (self.type) {
             inline .constant, .bin_op, .proc_call => |s| s.dest,
-            // inline .x => null,
         };
     }
 
@@ -85,17 +75,15 @@ pub const Instr = struct {
 pub const Block = struct {
     // arguments: []Reg,
     instrs: []Instr,
-    vreg_last_use: std.AutoHashMap(usize, std.ArrayList(VReg)),
+    vreg_last_use: std.AutoHashMap(VReg, usize),
 
     fn init(allocator: Allocator, instrs: []Instr) !Block {
-        var vreg_last_use: std.AutoHashMap(usize, std.ArrayList(VReg)) = .init(allocator);
+        var vreg_last_use: std.AutoHashMap(VReg, usize) = .init(allocator);
         for (0.., instrs) |i, instr| {
-            const kv = try vreg_last_use.getOrPut(i);
-            if (!kv.found_existing) kv.value_ptr.* = .init(allocator);
-            for (instr.sources().slice()) |src| {
-                if (std.mem.indexOfScalar(VReg, kv.value_ptr.items, src) == null)
-                    try kv.value_ptr.append(src);
-            }
+            for (instr.sources().slice()) |src|
+                try vreg_last_use.put(src, i);
+            if (instr.dest()) |dest|
+                try vreg_last_use.put(dest, i);
         }
         return .{
             .instrs = instrs,
@@ -134,10 +122,7 @@ const CompileContext = struct {
 
     fn compileStmt(self: *Self, stmt: parse.Stmt) !void {
         switch (stmt.type) {
-            .expr => |expr| try self.addInstr(.{
-                .loc = stmt.loc,
-                .type = .{ .discard = .{ .vreg = try self.compileExpr(expr) } },
-            }),
+            .expr => |expr| _ = try self.compileExpr(expr),
             .declare_var => |declare_var| {
                 const val = try self.compileExpr(declare_var.value);
                 const name = declare_var.ident.getIdent(self.source);
diff --git a/src/main.zig b/src/main.zig
index 15118d8..a9c77ac 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -62,6 +62,7 @@ pub fn main() !void {
     for (block.instrs) |instr| {
         std.debug.print("    {}\n", .{instr});
     }
+    std.debug.print("Last use of each virtual register: {}\n", .{fmtHashMap(block.vreg_last_use)});
     const elf = try codegen.create_elf(allocator, block);
     try output.writeAll(elf);
     std.debug.print("Run output:\n", .{});