store local variables on the stack

1 files changed, 18 insertions, 9 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index 97e160c..3fed244 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -332,7 +332,7 @@ const Instruction = union(enum) {
         fn init(opcode: Opcode, funct3: u3, rs1: Register, rs2: Register, imm: i12) Self {
             return .{ .s = .{
                 .opcode = opcode,
-                .imm4_0 = @intCast(imm & 0xf),
+                .imm4_0 = @intCast(imm & 0x1f),
                 .funct3 = funct3,
                 .rs1 = rs1,
                 .rs2 = rs2,
@@ -583,9 +583,10 @@ const Context = struct {
 
 const ProcedureContext = struct {
     register_allocator: RegisterAllocator(compile.VReg),
-    lvar_allocator: RegisterAllocator(compile.LVar),
     ctx: *Context,
     proc: compile.Procedure,
+    locals: std.AutoHashMap(compile.LVar, struct { stack_offset: i12 }),
+    stack_size: i12,
 
     // Current stuff that changes often, basically here to avoid prop drilling.
     block: ?*const compile.BasicBlock = null,
@@ -708,16 +709,16 @@ const ProcedureContext = struct {
     fn genAssignLocal(self: *Self, assign_local: compile.Instr.AssignLocal) !void {
         const src = self.register_allocator.get(assign_local.val);
         try self.freeUnusedVRegs();
-        const reg = try self.lvar_allocator.getOrAllocate(assign_local.local);
-        try self.emit(.addi(reg, src, 0));
+        std.log.debug("{}", .{self.locals.get(assign_local.local).?.stack_offset});
+        try self.emit(.sd(.sp, self.locals.get(assign_local.local).?.stack_offset, src));
     }
 
     fn genGetLocal(self: *Self, get_local: compile.Instr.GetLocal) !void {
         try self.freeUnusedVRegs();
 
-        const src = self.lvar_allocator.get(get_local.local);
         const reg = try self.register_allocator.allocate(get_local.dest);
-        try self.emit(.addi(reg, src, 0));
+
+        try self.emit(.ld(reg, .sp, self.locals.get(get_local.local).?.stack_offset));
     }
 
     fn codegenInstr(self: *Self, instr: compile.Instr) !void {
@@ -751,7 +752,7 @@ const ProcedureContext = struct {
     }
 
     fn prologue(self: *Self) !void {
-        try self.emit(.addi(.sp, .sp, -8));
+        try self.emit(.addi(.sp, .sp, -self.stack_size));
         try self.emit(.sd(.sp, 0, .ra));
 
         if (self.proc.param_reg) |reg| {
@@ -765,11 +766,18 @@ const ProcedureContext = struct {
             self.register_allocator.free(reg);
         }
         try self.emit(.ld(.ra, .sp, 0));
-        try self.emit(.addi(.sp, .sp, 8));
+        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 });
+        }
+        self.stack_size = @intCast(ptr);
+
         var first = true;
         for (self.proc.blocks) |block| {
             try self.ctx.block_addrs.putNoClobber(block.ref, self.ctx.instructions.items.len);
@@ -873,9 +881,10 @@ 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 }),
-            .lvar_allocator = try .init(allocator, &.{ .s11, .s10, .s9, .s8, .s7, .s6, .s5, .s4, .s3, .s2, .s1, .s0 }),
             .ctx = &ctx,
             .proc = proc,
+            .locals = .init(allocator),
+            .stack_size = undefined,
         };
         defer proc_ctx.deinit();