make local variables work separately from temporary values

3 files changed, 159 insertions, 215 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index d30026f..8c3f592 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -491,58 +491,66 @@ const Instruction = packed union {
     const Self = @This();
 };
 
-const RegisterAllocator = struct {
-    allocated: std.AutoHashMap(compile.VReg, Register),
-    available: std.ArrayList(Register),
-
-    fn init(allocator: Allocator) !RegisterAllocator {
-        var available: std.ArrayList(Register) = .init(allocator);
-        for ([_]Register{ .s11, .s10, .s9, .s8, .s7, .s6, .s5, .s4, .s3, .s2, .s1, .s0 }) |reg| {
-            try available.append(reg);
+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,
+            };
         }
-        var allocated: std.AutoHashMap(compile.VReg, Register) = .init(allocator);
-        try allocated.ensureTotalCapacity(@intCast(available.items.len));
-        return .{
-            .allocated = allocated,
-            .available = available,
-        };
-    }
 
-    fn deinit(self: *RegisterAllocator) void {
-        self.allocated.deinit();
-        self.available.deinit();
-    }
+        fn deinit(self: *Self) void {
+            self.allocated.deinit();
+            self.available.deinit();
+        }
 
-    fn get(self: *const RegisterAllocator, vreg: compile.VReg) Register {
-        return self.allocated.get(vreg).?;
-    }
+        fn get(self: *const Self, vreg: T) Register {
+            return self.allocated.get(vreg).?;
+        }
 
-    fn allocate(self: *RegisterAllocator, vreg: compile.VReg) !Register {
-        const reg = self.available.pop() orelse return error.OutOfRegisters;
-        self.allocated.putAssumeCapacityNoClobber(vreg, reg);
-        return reg;
-    }
+        fn getOrAllocate(self: *Self, vreg: T) !Register {
+            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 free(self: *RegisterAllocator, vreg: compile.VReg) void {
-        const ent = self.allocated.fetchRemove(vreg).?;
-        const reg = ent.value;
+        fn free(self: *Self, vreg: T) void {
+            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: *RegisterAllocator) !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: *RegisterAllocator, 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 Relocation = struct {
     instr: usize,
@@ -550,7 +558,8 @@ const Relocation = struct {
 };
 
 const Context = struct {
-    register_allocator: RegisterAllocator,
+    register_allocator: RegisterAllocator(compile.VReg),
+    lvar_allocator: RegisterAllocator(compile.LVar),
     instructions: std.ArrayList(Instruction),
     relocations: std.ArrayList(Relocation),
     block_starts: std.ArrayList(usize),
@@ -581,14 +590,14 @@ const Context = struct {
     /// collide with the sources. Should be called before allocating results to allow for more
     /// register re-use.
     fn freeUnusedVRegs(self: *Context) !void {
-        // TODO: make this do stuff again.
-        _ = self;
-        // 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.*);
-        //     }
-        // }
+        var it = self.register_allocator.allocated.keyIterator();
+        while (it.next()) |vreg| {
+            if (self.block.?.vreg_last_use.get(vreg.*)) |last_use| {
+                if (last_use <= self.current_instruction_index.?) {
+                    self.register_allocator.free(vreg.*);
+                }
+            }
+        }
     }
 
     fn genConstantInner(self: *Context, reg: Register, value: u64) !void {
@@ -729,16 +738,35 @@ const Context = struct {
     }
 
     fn genJump(self: *Context, jump: compile.Instr.Jump) !void {
+        try self.freeUnusedVRegs();
+
         try self.addRelocation(jump.to);
         try self.emit(.jal(.zero, 0));
     }
 
     fn genExit(self: *Context, _: compile.Instr.Exit) !void {
+        try self.freeUnusedVRegs();
+
         try self.emit(.addi(.a0, .zero, 0));
         try self.emit(.addi(.a7, .zero, 93));
         try self.emit(.ecall());
     }
 
+    fn genAssignLocal(self: *Context, 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));
+    }
+
+    fn genGetLocal(self: *Context, 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));
+    }
+
     fn codegenInstr(self: *Context, instr: compile.Instr) !void {
         switch (instr.type) {
             inline else => |ty| {
@@ -779,7 +807,8 @@ const Context = struct {
 
 pub fn create_elf(allocator: Allocator, proc: compile.Procedure) ![]u8 {
     var ctx: Context = .{
-        .register_allocator = try .init(allocator),
+        .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 }),
         .instructions = .init(allocator),
         .relocations = .init(allocator),
         .block_starts = .init(allocator),
@@ -809,7 +838,7 @@ pub fn create_elf(allocator: Allocator, proc: compile.Procedure) ![]u8 {
         }
     }
 
-    std.debug.print("allocated regs: {}\n", .{root.fmtHashMap(ctx.register_allocator.allocated)});
+    std.debug.assert(ctx.register_allocator.allocated.count() == 0);
 
     var output_buffer: std.ArrayList(u8) = .init(allocator);
     errdefer output_buffer.deinit();
diff --git a/src/compile.zig b/src/compile.zig
index 4d39293..93f240b 100644
--- a/src/compile.zig
+++ b/src/compile.zig
@@ -40,6 +40,8 @@ pub const Instr = struct {
         proc_call: ProcCall,
         branch: Branch,
         jump: Jump,
+        assign_local: AssignLocal,
+        get_local: GetLocal,
         exit: Exit,
     };
 
@@ -85,25 +87,41 @@ pub const Instr = struct {
 
     pub const Branch = struct {
         cond: VReg,
-        true: Jump,
-        false: Jump,
+        true: BlockRef,
+        false: BlockRef,
 
         pub const may_end_block = {};
 
         pub fn sources(_: Branch) Sources {
-            @panic("Sources not implemented for Branch!!!");
-            // return Sources.fromSlice(&.{self.cond}) catch unreachable;
+            return Sources.init(0) catch unreachable;
         }
     };
 
     pub const Jump = struct {
         to: BlockRef,
-        args: std.ArrayListUnmanaged(VReg) = .empty,
 
         pub const may_end_block = {};
 
-        pub fn sources(self: Jump) Sources {
-            return Sources.fromSlice(self.args) catch unreachable;
+        pub fn sources(_: Jump) Sources {
+            return Sources.init(0) catch unreachable;
+        }
+    };
+
+    pub const AssignLocal = struct {
+        local: LVar,
+        val: VReg,
+
+        pub fn sources(self: AssignLocal) Sources {
+            return Sources.fromSlice(&.{self.val}) catch unreachable;
+        }
+    };
+
+    pub const GetLocal = struct {
+        dest: VReg,
+        local: LVar,
+
+        pub fn sources(_: GetLocal) Sources {
+            return Sources.init(0) catch unreachable;
         }
     };
 
@@ -123,8 +141,8 @@ pub const Instr = struct {
 
     pub fn dest(self: *const Instr) ?VReg {
         return switch (self.type) {
-            inline .constant, .bin_op, .proc_call => |s| s.dest,
-            .branch, .jump, .exit => null,
+            inline .constant, .bin_op, .proc_call, .get_local => |s| s.dest,
+            .branch, .jump, .exit, .assign_local => null,
         };
     }
 
@@ -156,25 +174,20 @@ pub const Instr = struct {
                 },
             ),
             .branch => |branch| {
-                try writer.print("branch %{} ? ${}(", .{
+                try writer.print("branch %{} ? ${} : ${}", .{
                     @intFromEnum(branch.cond),
-                    @intFromEnum(branch.true.to),
+                    @intFromEnum(branch.true),
+                    @intFromEnum(branch.false),
                 });
-                for (branch.true.args.items, 0..) |arg, i| {
-                    try writer.print("{s}%{}", .{ if (i > 0) ", " else "", @intFromEnum(arg) });
-                }
-                try writer.print(") : ${}(", .{@intFromEnum(branch.false.to)});
-                for (branch.false.args.items, 0..) |arg, i| {
-                    try writer.print("{s}%{}", .{ if (i > 0) ", " else "", @intFromEnum(arg) });
-                }
-                try writer.print(")", .{});
             },
             .jump => |jump| {
-                try writer.print("jump ${}(", .{@intFromEnum(jump.to)});
-                for (jump.args.items, 0..) |arg, i| {
-                    try writer.print("{s}%{}", .{ if (i > 0) ", " else "", @intFromEnum(arg) });
-                }
-                try writer.print(")", .{});
+                try writer.print("jump ${}", .{@intFromEnum(jump.to)});
+            },
+            .assign_local => |assign_local| {
+                try writer.print("@{} = %{}", .{ @intFromEnum(assign_local.local), @intFromEnum(assign_local.val) });
+            },
+            .get_local => |get_local| {
+                try writer.print("%{} = @{}", .{ @intFromEnum(get_local.dest), @intFromEnum(get_local.local) });
             },
 
             .exit => |_| try writer.print("exit", .{}),
@@ -185,10 +198,10 @@ pub const Instr = struct {
 pub const Procedure = struct {
     blocks: []BasicBlock,
 
-    fn init(blocks: []BasicBlock) !Procedure {
-        // for (blocks) |*block| {
-        //     try block.finalize(allocator);
-        // }
+    fn init(allocator: Allocator, blocks: []BasicBlock) !Procedure {
+        for (blocks) |*block| {
+            try block.finalize(allocator);
+        }
 
         return .{ .blocks = blocks };
     }
@@ -202,13 +215,12 @@ pub const Procedure = struct {
 };
 
 pub const BasicBlock = struct {
-    params: std.AutoArrayHashMapUnmanaged(LVar, VReg) = .empty,
     instrs: std.ArrayListUnmanaged(Instr) = .empty,
 
     vreg_last_use: std.AutoHashMapUnmanaged(VReg, usize) = .empty,
 
     fn finalize(self: *BasicBlock, allocator: Allocator) !void {
-        std.debug.assert(self.instrs.items > 0);
+        std.debug.assert(self.instrs.items.len > 0);
         std.debug.assert(switch (self.instrs.getLast().type) {
             inline else => |ty| @hasDecl(@TypeOf(ty), "may_end_block"),
         });
@@ -232,24 +244,11 @@ pub const BasicBlock = struct {
 
     pub fn format(self: BasicBlock, comptime fmt: []const u8, options: std.fmt.FormatOptions, writer: anytype) !void {
         _ = .{ fmt, options };
-        try writer.print("(", .{});
-        {
-            var it = self.params.iterator();
-            var first = true;
-            while (it.next()) |ent| {
-                try writer.print("{s}%{}", .{ if (first) "" else ", ", @intFromEnum(ent.value_ptr.*) });
-                first = false;
-            }
-        }
-        try writer.print("):\n", .{});
+        try writer.print(":\n", .{});
         for (self.instrs.items) |instr| {
             try writer.print("    {}\n", .{instr});
         }
     }
-
-    const Info = struct {
-        local_current_vreg: std.AutoArrayHashMapUnmanaged(LVar, VReg) = .empty,
-    };
 };
 
 pub fn compile(allocator: Allocator, source: []const u8, block: parse.Block) !Procedure {
@@ -260,7 +259,6 @@ pub fn compile(allocator: Allocator, source: []const u8, block: parse.Block) !Pr
         .lvar_ctr = .init,
         .scope = .{ .locals = .empty, .parent = null },
         .blocks = .empty,
-        .block_infos = .empty,
         .current_block = @enumFromInt(0),
     };
     _ = try ctx.switchToNewBlock();
@@ -269,8 +267,7 @@ pub fn compile(allocator: Allocator, source: []const u8, block: parse.Block) !Pr
         .loc = .{ .start = 0, .end = 0 },
         .type = .{ .exit = .{} },
     });
-    try ctx.assignArguments();
-    return try .init(try ctx.blocks.toOwnedSlice(allocator));
+    return try .init(allocator, try ctx.blocks.toOwnedSlice(allocator));
 }
 
 const Scope = struct {
@@ -295,7 +292,6 @@ const CompileContext = struct {
     scope: Scope,
 
     blocks: std.ArrayListUnmanaged(BasicBlock),
-    block_infos: std.ArrayListUnmanaged(BasicBlock.Info),
     current_block: BlockRef,
 
     fn compileBlock(self: *CompileContext, block: parse.Block) !void {
@@ -334,7 +330,13 @@ const CompileContext = struct {
                 };
 
                 const val = try self.compileExpr(assign_var.value);
-                try self.assignLocalVar(local, val);
+                try self.addInstr(.{
+                    .loc = stmt.loc,
+                    .type = .{ .assign_local = .{
+                        .local = local,
+                        .val = val,
+                    } },
+                });
             },
             .@"while" => |@"while"| {
                 const curr = self.current_block;
@@ -352,8 +354,8 @@ const CompileContext = struct {
                     .loc = stmt.loc,
                     .type = .{ .branch = .{
                         .cond = cond,
-                        .true = .{ .to = do },
-                        .false = .{ .to = after },
+                        .true = do,
+                        .false = after,
                     } },
                 });
                 self.current_block = do;
@@ -409,12 +411,21 @@ const CompileContext = struct {
             },
             .identifier => {
                 var scope: ?*Scope = &self.scope;
-                while (scope) |s| : (scope = s.parent) {
-                    if (s.locals.get(expr.loc.getIdent(self.source))) |local| {
-                        return self.getLocalVar(local);
+                const local: LVar = blk: {
+                    while (scope) |s| : (scope = s.parent) {
+                        if (s.locals.get(expr.loc.getIdent(self.source))) |local| {
+                            break :blk local;
+                        }
                     }
-                }
-                return error.UnknownVariable;
+                    return error.UnknownVariable;
+                };
+                try self.addInstr(.{
+                    .loc = expr.loc,
+                    .type = .{ .get_local = .{
+                        .dest = dest,
+                        .local = local,
+                    } },
+                });
             },
             .@"if" => |@"if"| {
                 const cond = try self.compileExpr(@"if".cond);
@@ -444,8 +455,8 @@ const CompileContext = struct {
                     .loc = expr.loc,
                     .type = .{ .branch = .{
                         .cond = cond,
-                        .true = .{ .to = t },
-                        .false = .{ .to = f },
+                        .true = t,
+                        .false = f,
                     } },
                 });
 
@@ -455,83 +466,9 @@ const CompileContext = struct {
         return dest;
     }
 
-    fn assignArguments(self: *CompileContext) !void {
-        var immediate_predecessors: std.AutoHashMapUnmanaged(usize, std.ArrayListUnmanaged(usize)) = .empty;
-        defer immediate_predecessors.deinit(self.allocator);
-        defer {
-            var it = immediate_predecessors.valueIterator();
-            while (it.next()) |val| val.deinit(self.allocator);
-        }
-
-        for (self.blocks.items, 0..) |block, i| {
-            switch (block.instrs.getLast().type) {
-                .branch => |branch| {
-                    const pt = try immediate_predecessors.getOrPut(self.allocator, @intFromEnum(branch.true.to));
-                    if (!pt.found_existing) {
-                        pt.value_ptr.* = .empty;
-                    }
-                    try pt.value_ptr.append(self.allocator, i);
-
-                    const pf = try immediate_predecessors.getOrPut(self.allocator, @intFromEnum(branch.false.to));
-                    if (!pf.found_existing) {
-                        pf.value_ptr.* = .empty;
-                    }
-                    try pf.value_ptr.append(self.allocator, i);
-                },
-                .jump => |jump| {
-                    const p = try immediate_predecessors.getOrPut(self.allocator, @intFromEnum(jump.to));
-                    if (!p.found_existing) {
-                        p.value_ptr.* = .empty;
-                    }
-                    try p.value_ptr.append(self.allocator, i);
-                },
-                else => {},
-            }
-        }
-
-        var block_ids: std.ArrayListUnmanaged(usize) = .empty;
-        for (0..self.blocks.items.len) |i| try block_ids.append(self.allocator, i);
-        while (block_ids.pop()) |ref| {
-            const block = &self.blocks.items[ref];
-            const info = &self.block_infos.items[ref];
-            const ty = &block.instrs.items[block.instrs.items.len - 1].type;
-            const got_new_params = switch (ty.*) {
-                .branch => blk: {
-                    const t = try self.assignArgumentsInner(block, info, &ty.branch.true);
-                    const f = try self.assignArgumentsInner(block, info, &ty.branch.false);
-                    break :blk t or f;
-                },
-                .jump => try self.assignArgumentsInner(block, info, &ty.jump),
-                else => false,
-            };
-            if (got_new_params) {
-                try block_ids.appendSlice(self.allocator, immediate_predecessors.get(ref).?.items);
-            }
-        }
-    }
-
-    fn assignArgumentsInner(self: *CompileContext, block: *BasicBlock, info: *BasicBlock.Info, jump: *Instr.Jump) !bool {
-        var got_new_params = false;
-        var it = self.blocks.items[@intFromEnum(jump.to)].params.iterator();
-        // Skip over the arguments that we've already handled in an earlier call
-        for (0..jump.args.items.len) |_| _ = it.next();
-        while (it.next()) |ent| {
-            try jump.args.append(self.allocator, info.local_current_vreg.get(ent.key_ptr.*) orelse blk: {
-                const reg = self.vreg_ctr.get();
-                try block.params.put(self.allocator, ent.key_ptr.*, reg);
-                try info.local_current_vreg.put(self.allocator, ent.key_ptr.*, reg);
-                got_new_params = true;
-                break :blk reg;
-            });
-        }
-
-        return got_new_params;
-    }
-
     fn switchToNewBlock(self: *CompileContext) !BlockRef {
         const ref: BlockRef = @enumFromInt(self.blocks.items.len);
         try self.blocks.append(self.allocator, .{});
-        try self.block_infos.append(self.allocator, .{});
         self.current_block = ref;
         return ref;
     }
@@ -540,22 +477,4 @@ const CompileContext = struct {
         try self.blocks.items[@intFromEnum(self.current_block)]
             .instrs.append(self.allocator, instr);
     }
-
-    fn getLocalVar(self: *CompileContext, local: LVar) !VReg {
-        if (self.block_infos.items[@intFromEnum(self.current_block)]
-            .local_current_vreg.get(local)) |vreg| return vreg;
-
-        const vreg = self.vreg_ctr.get();
-        try self.blocks.items[@intFromEnum(self.current_block)]
-            .params.put(self.allocator, local, vreg);
-        try self.block_infos.items[@intFromEnum(self.current_block)]
-            .local_current_vreg.put(self.allocator, local, vreg);
-
-        return vreg;
-    }
-
-    fn assignLocalVar(self: *CompileContext, local: LVar, vreg: VReg) !void {
-        try self.block_infos.items[@intFromEnum(self.current_block)]
-            .local_current_vreg.put(self.allocator, local, vreg);
-    }
 };
diff --git a/src/main.zig b/src/main.zig
index becbf97..5ee36da 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -67,15 +67,8 @@ pub fn main() !void {
     }
     const procedure = try compile.compile(allocator, source, ast);
     std.debug.print("Bytecode instructions:\n{}", .{procedure});
-    // for (procedure.blocks, 0..) |block, i| {
-    //     std.debug.print("  ${}:\n", .{i});
-    //     for (block.instrs.items) |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, procedure);
-    // try out_file.writer().writeAll(elf);
+    const elf = try codegen.create_elf(allocator, procedure);
+    try out_file.writer().writeAll(elf);
     std.debug.print("Run output:\n", .{});
     if (run) {
         out_file.close();
@@ -104,7 +97,10 @@ fn HashMapFormatter(HashMap: type) type {
             var it = hash_map.iterator();
             var first = true;
             while (it.next()) |kv| {
-                try writer.print("{s} {any}: {any}", .{ if (first) "" else ",", kv.key_ptr.*, kv.value_ptr.* });
+                try writer.print(
+                    "{s} {" ++ (if (@TypeOf(kv.key_ptr.*) == []const u8) "s" else "any") ++ "}: {any}",
+                    .{ if (first) "" else ",", kv.key_ptr.*, kv.value_ptr.* },
+                );
                 first = false;
             }
             try writer.writeAll(if (first) "}" else " }");