1 files changed, 16 insertions, 10 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index a1c2171..00b3857 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -563,7 +563,7 @@ const Relocation = struct {
 const Context = struct {
     instructions: std.ArrayList(Instruction),
     relocations: std.ArrayList(Relocation),
-    block_starts: std.ArrayList(usize),
+    block_addrs: std.AutoHashMap(compile.BlockRef, usize),
     print_block: compile.BlockRef,
     read_int_block: compile.BlockRef,
 
@@ -767,7 +767,7 @@ const ProcedureContext = struct {
 
     fn codegenProc(self: *Self, proc: compile.Procedure) !void {
         for (proc.blocks) |block| {
-            try self.ctx.block_starts.append(self.ctx.instructions.items.len);
+            try self.ctx.block_addrs.putNoClobber(block.ref, self.ctx.instructions.items.len);
             try self.codegenBlock(block);
         }
     }
@@ -845,17 +845,23 @@ fn codegenReadInt(self: *Context) !void {
     try self.emit(.jalr(.zero, .ra, 0));
 }
 
-pub fn create_elf(allocator: Allocator, proc: compile.Procedure) ![]u8 {
+pub fn create_elf(allocator: Allocator, mod: compile.Module) ![]u8 {
+    var maxBlockRef: usize = 0;
+    for (mod.procedures) |proc| {
+        for (proc.blocks) |block| {
+            maxBlockRef = @max(maxBlockRef, @intFromEnum(block.ref));
+        }
+    }
     var ctx: Context = .{
         .instructions = .init(allocator),
         .relocations = .init(allocator),
-        .block_starts = .init(allocator),
-        .print_block = @enumFromInt(proc.blocks.len),
-        .read_int_block = @enumFromInt(proc.blocks.len + 1),
+        .block_addrs = .init(allocator),
+        .print_block = @enumFromInt(maxBlockRef + 1),
+        .read_int_block = @enumFromInt(maxBlockRef + 2),
     };
     defer ctx.deinit();
 
-    {
+    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 }),
@@ -867,15 +873,15 @@ pub fn create_elf(allocator: Allocator, proc: compile.Procedure) ![]u8 {
         std.debug.assert(proc_ctx.register_allocator.allocated.count() == 0);
     }
 
-    try ctx.block_starts.append(ctx.instructions.items.len);
+    try ctx.block_addrs.putNoClobber(ctx.print_block, ctx.instructions.items.len);
     try codegenPrint(&ctx);
-    try ctx.block_starts.append(ctx.instructions.items.len);
+    try ctx.block_addrs.putNoClobber(ctx.read_int_block, ctx.instructions.items.len);
     try codegenReadInt(&ctx);
 
     // TODO: make this less sheiße
     for (ctx.relocations.items) |relocation| {
         const instr = &ctx.instructions.items[relocation.instr];
-        const target: isize = @intCast(ctx.block_starts.items[@intFromEnum(relocation.target)]);
+        const target: isize = @intCast(ctx.block_addrs.get(relocation.target).?);
         const from: isize = @intCast(relocation.instr);
         switch (instr.*) {
             .j => |j| {