use saved rregisters when needed

& store and restore saved registers on the stack

2 files changed, 106 insertions, 59 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index bf84412..8aae214 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -494,66 +494,74 @@ const Instruction = union(enum) {
     const Self = @This();
 };
 
-fn RegisterAllocator(T: type) type {
-    return struct {
-        allocated: std.AutoHashMap(T, Register),
-        available: std.ArrayList(Register),
-
-        fn init(allocator: Allocator, regs: []const Register) !Self {
-            var available: std.ArrayList(Register) = .init(allocator);
-            try available.appendSlice(regs);
-            var allocated: std.AutoHashMap(T, Register) = .init(allocator);
-            try allocated.ensureTotalCapacity(@intCast(available.items.len));
-            return .{
-                .allocated = allocated,
-                .available = available,
-            };
-        }
+const RegisterAllocator = struct {
+    allocated: std.AutoHashMap(compile.VReg, Register),
+    available: std.ArrayList(Register),
+    to_ignore: std.AutoHashMap(compile.VReg, SavedRegInfo),
+
+    fn init(
+        allocator: Allocator,
+        to_ignore: std.AutoHashMap(compile.VReg, SavedRegInfo),
+    ) !Self {
+        var available: std.ArrayList(Register) = .init(allocator);
+        try available.appendSlice(&.{ .t6, .t5, .t4, .t3, .t2, .t1, .t0 });
+        var allocated: std.AutoHashMap(compile.VReg, Register) = .init(allocator);
+        try allocated.ensureTotalCapacity(@intCast(available.items.len));
+
+        return .{
+            .allocated = allocated,
+            .available = available,
+            .to_ignore = to_ignore,
+        };
+    }
 
-        fn deinit(self: *Self) void {
-            self.allocated.deinit();
-            self.available.deinit();
-        }
+    fn deinit(self: *Self) void {
+        self.allocated.deinit();
+        self.available.deinit();
+    }
 
-        fn get(self: *const Self, vreg: T) Register {
-            return self.allocated.get(vreg).?;
-        }
+    fn get(self: *const Self, vreg: compile.VReg) Register {
+        if (self.to_ignore.get(vreg)) |info| return info.reg;
+        return self.allocated.get(vreg).?;
+    }
 
-        fn getOrAllocate(self: *Self, vreg: T) !Register {
-            const reg = self.allocated.get(vreg) orelse
-                try self.allocate(vreg);
-            return reg;
-        }
+    fn getOrAllocate(self: *Self, vreg: compile.VReg) !Register {
+        if (self.to_ignore.get(vreg)) |info| return info.reg;
+        const reg = self.allocated.get(vreg) orelse
+            try self.allocate(vreg);
+        return reg;
+    }
 
-        fn allocate(self: *Self, vreg: T) !Register {
-            const reg = self.available.pop() orelse return error.OutOfRegisters;
-            self.allocated.putAssumeCapacityNoClobber(vreg, reg);
-            return reg;
-        }
+    fn allocate(self: *Self, vreg: compile.VReg) !Register {
+        if (self.to_ignore.get(vreg)) |info| return info.reg;
+        const reg = self.available.pop() orelse return error.OutOfRegisters;
+        self.allocated.putAssumeCapacityNoClobber(vreg, reg);
+        return reg;
+    }
 
-        fn free(self: *Self, vreg: T) void {
-            const ent = self.allocated.fetchRemove(vreg).?;
-            const reg = ent.value;
+    fn free(self: *Self, vreg: compile.VReg) void {
+        if (self.to_ignore.get(vreg) != null) return;
+        const ent = self.allocated.fetchRemove(vreg).?;
+        const reg = ent.value;
 
-            std.debug.assert(std.mem.indexOfScalar(Register, self.available.items, reg) == null);
-            return self.available.appendAssumeCapacity(reg);
-        }
+        std.debug.assert(std.mem.indexOfScalar(Register, self.available.items, reg) == null);
+        return self.available.appendAssumeCapacity(reg);
+    }
 
-        fn allocateAux(self: *Self) !Register {
-            const reg = self.available.pop() orelse return error.OutOfRegisters;
-            return reg;
-        }
+    fn allocateAux(self: *Self) !Register {
+        const reg = self.available.pop() orelse return error.OutOfRegisters;
+        return reg;
+    }
 
-        fn freeAux(self: *Self, reg: Register) void {
-            var it = self.allocated.valueIterator();
-            while (it.next()) |r| std.debug.assert(reg != r.*);
-            std.debug.assert(std.mem.indexOfScalar(Register, self.available.items, reg) == null);
-            return self.available.appendAssumeCapacity(reg);
-        }
+    fn freeAux(self: *Self, reg: Register) void {
+        var it = self.allocated.valueIterator();
+        while (it.next()) |r| std.debug.assert(reg != r.*);
+        std.debug.assert(std.mem.indexOfScalar(Register, self.available.items, reg) == null);
+        return self.available.appendAssumeCapacity(reg);
+    }
 
-        const Self = @This();
-    };
-}
+    const Self = @This();
+};
 
 const Relocation = struct {
     instr: usize,
@@ -561,6 +569,7 @@ const Relocation = struct {
 };
 
 const Context = struct {
+    allocator: std.mem.Allocator,
     instructions: std.ArrayList(Instruction),
     relocations: std.ArrayList(Relocation),
     block_addrs: std.AutoHashMap(compile.BlockRef, usize),
@@ -581,10 +590,16 @@ const Context = struct {
     }
 };
 
+const SavedRegInfo = struct {
+    stored_stack_offset: i12,
+    reg: Register,
+};
+
 const ProcedureContext = struct {
-    register_allocator: RegisterAllocator(compile.VReg),
+    register_allocator: RegisterAllocator,
     ctx: *Context,
     proc: compile.Procedure,
+    saved_values: std.AutoHashMap(compile.VReg, SavedRegInfo),
     locals: std.AutoHashMap(compile.LVar, struct { stack_offset: i12 }),
     stack_size: i12,
 
@@ -761,21 +776,43 @@ const ProcedureContext = struct {
         if (self.proc.param_lvar) |lvar| {
             try self.emit(.sd(.sp, self.locals.get(lvar).?.stack_offset, .a0));
         }
+        var it = self.saved_values.valueIterator();
+        while (it.next()) |info| {
+            try self.emit(.sd(.sp, info.stored_stack_offset, info.reg));
+        }
     }
 
     fn epilogue(self: *Self) !void {
         try self.emit(.ld(.ra, .sp, 0));
+
+        var it = self.saved_values.valueIterator();
+        while (it.next()) |info| {
+            try self.emit(.ld(info.reg, .sp, info.stored_stack_offset));
+        }
+
         try self.emit(.addi(.sp, .sp, self.stack_size));
         try self.emit(.jalr(.zero, .ra, 0));
     }
 
     fn codegenProc(self: *Self) !void {
-        var it = self.proc.locals.keyIterator();
         var ptr: i12 = 8; // return address is at [0,8)
-        while (it.next()) |lvar| : (ptr += 8) {
-            try self.locals.putNoClobber(lvar.*, .{ .stack_offset = ptr });
+        {
+            var it = self.proc.locals.keyIterator();
+            while (it.next()) |lvar| : (ptr += 8) {
+                try self.locals.putNoClobber(lvar.*, .{ .stack_offset = ptr });
+            }
+        }
+        var s_regs: []const Register = &.{ .s0, .s1, .s2, .s3, .s3, .s4, .s5, .s6, .s7, .s8, .s9, .s10, .s11 };
+        for (self.proc.blocks) |block| {
+            var it = block.vreg_used_during_call.keyIterator();
+            while (it.next()) |vreg| {
+                try self.saved_values.putNoClobber(vreg.*, .{ .stored_stack_offset = ptr, .reg = s_regs[0] });
+                ptr += 8;
+                s_regs = s_regs[1..];
+            }
         }
         self.stack_size = @intCast(ptr);
+        self.register_allocator = try .init(self.ctx.allocator, self.saved_values);
 
         var first = true;
         for (self.proc.blocks) |block| {
@@ -871,6 +908,7 @@ fn codegenExit(self: *Context) !void {
 
 pub fn create_elf(allocator: Allocator, mod: compile.Module) ![]u8 {
     var ctx: Context = .{
+        .allocator = allocator,
         .instructions = .init(allocator),
         .relocations = .init(allocator),
         .block_addrs = .init(allocator),
@@ -879,10 +917,11 @@ pub fn create_elf(allocator: Allocator, mod: compile.Module) ![]u8 {
 
     for (mod.procedures) |proc| {
         var proc_ctx: ProcedureContext = .{
-            .register_allocator = try .init(allocator, &.{ .t6, .t5, .t4, .t3, .t2, .t1, .t0 }),
+            .register_allocator = undefined,
             .ctx = &ctx,
             .proc = proc,
             .locals = .init(allocator),
+            .saved_values = .init(allocator),
             .stack_size = undefined,
         };
         defer proc_ctx.deinit();
diff --git a/src/compile.zig b/src/compile.zig
index 0ce170c..3ee3eb0 100644
--- a/src/compile.zig
+++ b/src/compile.zig
@@ -94,6 +94,7 @@ pub const BasicBlock = struct {
     instrs: std.ArrayListUnmanaged(Instr) = .empty,
 
     vreg_last_use: std.AutoHashMapUnmanaged(VReg, usize) = .empty,
+    vreg_used_during_call: std.AutoHashMapUnmanaged(VReg, void) = .empty,
 
     fn finalize(self: *BasicBlock, allocator: Allocator) !void {
         std.debug.assert(self.instrs.items.len > 0);
@@ -102,11 +103,18 @@ pub const BasicBlock = struct {
         });
 
         self.vreg_last_use = .empty;
+        var set_at: std.AutoHashMapUnmanaged(VReg, usize) = .empty;
+        var last_call: usize = 0;
         for (0.., self.instrs.items) |i, instr| {
-            for (instr.sources().slice()) |src|
-                try self.vreg_last_use.put(allocator, src, i);
-            if (instr.dest()) |dest|
+            if (instr.dest()) |dest| {
                 try self.vreg_last_use.put(allocator, dest, i);
+                try set_at.put(allocator, dest, i);
+            }
+            for (instr.sources().slice()) |src| {
+                try self.vreg_last_use.put(allocator, src, i);
+                if (set_at.get(src).? < last_call) try self.vreg_used_during_call.put(allocator, src, {});
+            }
+            if (instr.type == .call) last_call = i;
         }
     }