1 files changed, 16 insertions, 8 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index 1e8ed41..fe58f74 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -560,10 +560,11 @@ const Context = struct {
         try self.instructions.append(inst);
     }
 
-    fn maybeFreeSources(self: *Context, vregs: compile.Instr.Sources) !void {
-        for (vregs.slice()) |src| {
-            if (self.block.?.vreg_last_use.get(src) == self.current_instruction_index.?) {
-                self.register_allocator.free(src);
+    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);
             }
         }
     }
@@ -601,7 +602,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(bin_op.sources());
+        try self.maybeFreeSources();
         const reg = try self.register_allocator.allocate(bin_op.dest);
         switch (bin_op.op) {
             .add => try self.emit(.add(reg, lhs, rhs)),
@@ -616,10 +617,9 @@ const Context = struct {
         defer self.register_allocator.freeAux(quot);
         const digit = try self.register_allocator.allocateAux();
         defer self.register_allocator.freeAux(digit);
-        const count = try self.register_allocator.allocateAux();
-        defer self.register_allocator.freeAux(count);
+        const count = try self.register_allocator.allocate(print.dest);
 
-        try self.maybeFreeSources(print.sources());
+        try self.maybeFreeSources();
 
         try self.emit(.addi(digit, .zero, '\n'));
         try self.emit(.addi(.sp, .sp, -1));
@@ -642,6 +642,12 @@ const Context = struct {
         try self.emit(.addi(.a2, count, 0)); // count = count
         try self.emit(.ecall()); // syscall(no, fd, buf, count)
         try self.emit(.add(.sp, .sp, count));
+
+        try self.emit(.addi(count, count, -1));
+    }
+
+    fn genDiscard(self: *Context, _: compile.Instr.Discard) !void {
+        try self.maybeFreeSources();
     }
 
     fn codegenInstr(self: *Context, instr: compile.Instr) !void {
@@ -687,6 +693,8 @@ pub fn create_elf(allocator: Allocator, block: compile.Block) ![]u8 {
         .ecall(),
     });
 
+    std.debug.print("allocated regs: {}\n", .{root.fmtHashMap(ctx.register_allocator.allocated)});
+
     var output_buffer: std.ArrayList(u8) = .init(allocator);
     errdefer output_buffer.deinit();
     try output_buffer.appendNTimes(undefined, @sizeOf(elf.Elf64_Ehdr) + @sizeOf(elf.Elf64_Phdr));