store local variables on the stack

5 files changed, 61 insertions, 17 deletions

diff --git a/fibonacci.hgn b/fibonacci.hgn
index b1ce582..2a0cd72 100644
--- a/fibonacci.hgn
+++ b/fibonacci.hgn
@@ -13,6 +13,11 @@ fib := proc(n) {
         a = b
         b = c
         i = i + 1
+        # p(i)
     }
     print(a)
 }
+
+p := proc(x) {
+    x2 := x
+}
diff --git a/src/Lexer.zig b/src/Lexer.zig
index 4a4d895..19e9cd4 100644
--- a/src/Lexer.zig
+++ b/src/Lexer.zig
@@ -134,7 +134,7 @@ fn integerLiteral(self: *Self) Token {
 fn identifierOrKeyword(self: *Self) Token {
     while (true) {
         const c = self.peekChar() orelse 0;
-        if ('a' <= c and c <= 'z' or 'A' <= c and c <= 'Z' or c == '_') {
+        if ('a' <= c and c <= 'z' or 'A' <= c and c <= 'Z' or c == '_' or '0' <= c and c <= '9') {
             _ = self.eatChar();
             continue;
         }
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();
 
diff --git a/src/compile.zig b/src/compile.zig
index 80a9a17..636cf01 100644
--- a/src/compile.zig
+++ b/src/compile.zig
@@ -56,14 +56,26 @@ pub const Procedure = struct {
     name: []const u8,
     blocks: []BasicBlock,
     param_reg: ?VReg,
-
-    fn init(allocator: Allocator, name: []const u8, blocks: []BasicBlock, param_reg: ?VReg) !Procedure {
+    locals: std.AutoHashMap(LVar, void),
+
+    fn init(
+        allocator: Allocator,
+        name: []const u8,
+        blocks: []BasicBlock,
+        param_reg: ?VReg,
+        locals: std.AutoHashMap(LVar, void),
+    ) !Procedure {
         for (blocks) |*block| {
             try block.finalize(allocator);
             if (param_reg) |p| _ = block.vreg_last_use.remove(p);
         }
 
-        return .{ .name = name, .blocks = blocks, .param_reg = param_reg };
+        return .{
+            .name = name,
+            .blocks = blocks,
+            .param_reg = param_reg,
+            .locals = locals,
+        };
     }
 
     pub fn format(self: Procedure, comptime fmt: []const u8, options: std.fmt.FormatOptions, writer: anytype) !void {
@@ -352,6 +364,7 @@ fn compileProcedure(
         .vreg_ctr = .init,
         .lvar_ctr = .init,
         .scope = .{ .locals = .empty, .parent = null },
+        .locals = .empty,
         .param = .{ .name = "", .reg = undefined },
         .blocks = try .init(ctx.allocator, &.{first_block}, &.{.{ .ref = first_block }}),
         .current_block = first_block,
@@ -379,7 +392,13 @@ fn compileProcedure(
         std.debug.assert(kv.key_ptr.* == kv.value_ptr.ref);
         blocks[i] = kv.value_ptr.*;
     }
-    return try .init(ctx.allocator, name.getIdent(ctx.source), blocks, param_reg);
+    return try .init(
+        ctx.allocator,
+        name.getIdent(ctx.source),
+        blocks,
+        param_reg,
+        pctx.locals.promote(ctx.allocator),
+    );
 }
 
 const ProcedureContext = struct {
@@ -389,6 +408,7 @@ const ProcedureContext = struct {
     lvar_ctr: IdCounter(LVar),
 
     scope: Scope,
+    locals: std.AutoHashMapUnmanaged(LVar, void),
     param: struct { name: []const u8, reg: VReg },
 
     blocks: std.AutoArrayHashMapUnmanaged(BlockRef, BasicBlock),
@@ -420,6 +440,7 @@ const ProcedureContext = struct {
                     const local = self.lvar_ctr.get();
                     const name = assign_var.ident.getIdent(self.ctx.source);
                     try self.scope.locals.put(self.ctx.allocator, name, local);
+                    try self.locals.put(self.ctx.allocator, local, {});
                     break :blk local;
                 } else blk: {
                     var scope: ?*Scope = &self.scope;
diff --git a/src/main.zig b/src/main.zig
index 55b7e0e..d0e94db 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -74,15 +74,24 @@ fn HashMapFormatter(HashMap: type) type {
             options: std.fmt.FormatOptions,
             writer: anytype,
         ) !void {
-            _ = fmt;
+            const k_fmt, const v_fmt = comptime blk: {
+                var fmt_it = std.mem.splitScalar(u8, fmt, ',');
+                const k_fmt = fmt_it.next() orelse "";
+                const v_fmt = fmt_it.next() orelse "";
+                break :blk .{ k_fmt, v_fmt };
+            };
             _ = options;
             try writer.writeAll("{");
             var it = hash_map.iterator();
             var first = true;
             while (it.next()) |kv| {
                 try writer.print(
-                    "{s} {" ++ (if (@TypeOf(kv.key_ptr.*) == []const u8) "s" else "any") ++ "}: {any}",
-                    .{ if (first) "" else ",", kv.key_ptr.*, kv.value_ptr.* },
+                    "{s} {" ++ k_fmt ++ "}: {" ++ v_fmt ++ "}",
+                    .{
+                        if (first) "" else ",",
+                        kv.key_ptr.*,
+                        kv.value_ptr.*,
+                    },
                 );
                 first = false;
             }