1 files changed, 74 insertions, 7 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index 89daaa4..59b7040 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -2,7 +2,7 @@ const std = @import("std");
 const elf = std.elf;
 const Allocator = std.mem.Allocator;
 const root = @import("root");
-const Block = root.Block;
+const compile = root.compile;
 
 const Register = enum(u5) {
     // zig fmt: off
@@ -304,7 +304,7 @@ const Instruction = packed union {
         fn init(opcode: Opcode, funct3: u3, rs1: Register, rs2: Register, imm: i13) Self {
             std.debug.assert(imm % 2 == 0);
             const umm: u13 = @bitCast(imm);
-            return .{ .s = .{
+            return .{ .b = .{
                 .opcode = opcode,
                 .imm11 = @truncate(umm >> 11),
                 .imm4_1 = @truncate(umm >> 1),
@@ -371,7 +371,7 @@ const Instruction = packed union {
         imm12_31: u20,
 
         fn init(opcode: Opcode, rd: Register, imm: i20) Self {
-            return .{ .s = .{
+            return .{ .u = .{
                 .opcode = opcode,
                 .rd = rd,
                 .imm12_31 = @bitCast(imm),
@@ -419,17 +419,84 @@ const Instruction = packed union {
     const Self = @This();
 };
 
-pub fn create_elf(allocator: Allocator, block: Block) ![]u8 {
-    _ = block;
+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{ .t6, .t5, .t4, .t3, .t2, .t1, .t0 }) |reg| {
+            try available.append(reg);
+        }
+        var allocated: std.AutoHashMap(compile.VReg, Register) = .init(allocator);
+        try allocated.ensureTotalCapacity(@intCast(available.items.len));
+        return .{
+            .allocated = allocated,
+            .available = available,
+        };
+    }
+
+    fn get(self: *const RegisterAllocator, vreg: compile.VReg) Register {
+        return self.allocated.get(vreg).?;
+    }
+
+    fn allocate(self: *RegisterAllocator, vreg: compile.VReg) ?Register {
+        const reg = self.available.pop() orelse return null;
+        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;
 
+        std.debug.assert(std.mem.indexOfScalar(Register, self.available.items, reg) == null);
+        return self.available.appendAssumeCapacity(reg);
+    }
+};
+
+pub fn create_elf(allocator: Allocator, block: compile.Block) ![]u8 {
     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));
+
     const output = output_buffer.writer();
+    var register_allocator: RegisterAllocator = try .init(allocator);
+
+    for (0.., block.instrs) |i, instr| {
+        switch (instr.type) {
+            .constant => |constant| {
+                const reg = register_allocator.allocate(constant.dest) orelse return error.OutOfRegisters;
+
+                if (constant.value <= std.math.maxInt(i12)) {
+                    try output.writeInt(u32, @bitCast(Instruction.addi(reg, .zero, @intCast(constant.value))), .little);
+                } else if (constant.value <= std.math.maxInt(i32)) {
+                    // If the higest bit in the immediate in addi is set, it will be sign extended. We negate that by adding one more to the immediate for lui.
+                    try output.writeInt(u32, @bitCast(Instruction.lui(reg, @intCast((constant.value >> 12) + if (constant.value & (1 << 11) != 0) @as(u64, 1) else 0))), .little);
+                    try output.writeInt(u32, @bitCast(Instruction.addi(reg, reg, @bitCast(@as(u12, @truncate(constant.value))))), .little);
+                } else {
+                    unreachable; // TODO
+                }
+            },
+            .bin_op => |bin_op| {
+                const lhs = register_allocator.get(bin_op.lhs);
+                const rhs = register_allocator.get(bin_op.rhs);
+                for (instr.sources().slice()) |src| {
+                    if (block.vreg_last_use.get(src) == i) {
+                        register_allocator.free(src);
+                    }
+                }
+                const reg = register_allocator.allocate(bin_op.dest) orelse return error.OutOfRegisters;
+                switch (bin_op.op) {
+                    .add => try output.writeInt(u32, @bitCast(Instruction.add(reg, lhs, rhs)), .little),
+                }
+            },
+        }
+    }
+
     for ([_]Instruction{
+        .addi(.a0, register_allocator.get(block.instrs[block.instrs.len - 1].dest()), 0),
         .addi(.a7, .zero, 93),
-        .addi(.a0, .zero, 71),
-        .xori(.a0, .a0, 2),
         .ecall(),
     }) |instr| {
         try output.writeInt(u32, @bitCast(instr), .little);