1 files changed, 40 insertions, 38 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index 00b3857..0e70a0b 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -564,8 +564,6 @@ const Context = struct {
     instructions: std.ArrayList(Instruction),
     relocations: std.ArrayList(Relocation),
     block_addrs: std.AutoHashMap(compile.BlockRef, usize),
-    print_block: compile.BlockRef,
-    read_int_block: compile.BlockRef,
 
     fn addRelocation(self: *Context, target: compile.BlockRef) !void {
         try self.relocations.append(.{
@@ -667,29 +665,16 @@ const ProcedureContext = struct {
     }
 
     fn genProcCall(self: *Self, call: compile.Instr.ProcCall) !void {
-        switch (call.proc) {
-            .print => {
-                const arg = self.register_allocator.get(call.arg);
+        const arg = self.register_allocator.get(call.arg);
 
-                try self.freeUnusedVRegs();
-
-                const result = try self.register_allocator.allocate(call.dest);
-
-                try self.emit(.addi(.a0, arg, 0));
-                try self.ctx.addRelocation(self.ctx.print_block);
-                try self.emit(.jal(.ra, 0));
-                try self.emit(.addi(result, .a0, 0));
-            },
-            .read_int => {
-                try self.freeUnusedVRegs();
+        try self.freeUnusedVRegs();
 
-                const result = try self.register_allocator.allocate(call.dest);
+        const result = try self.register_allocator.allocate(call.dest);
 
-                try self.ctx.addRelocation(self.ctx.read_int_block);
-                try self.emit(.jal(.ra, 0));
-                try self.emit(.addi(result, .a0, 0));
-            },
-        }
+        try self.emit(.addi(.a0, arg, 0));
+        try self.ctx.addRelocation(call.proc);
+        try self.emit(.jal(.ra, 0));
+        try self.emit(.addi(result, .a0, 0));
     }
 
     fn genBranch(self: *Self, branch: compile.Instr.Branch) !void {
@@ -710,14 +695,13 @@ const ProcedureContext = struct {
         try self.emit(.jal(.zero, 0));
     }
 
-    fn genExit(self: *Self, _: compile.Instr.Exit) !void {
+    fn genReturn(self: *Self, ret: compile.Instr.Return) !void {
+        const val = self.register_allocator.get(ret.val);
+
         try self.freeUnusedVRegs();
 
-        try self.emit(.addi(.a0, .zero, 0));
-        try self.emit(.addi(.a7, .zero, @intFromEnum(std.os.linux.syscalls.RiscV64.exit)));
-        try self.emit(.ecall());
-        // This will never be run, but makes binary ninja understand that this exits the function.
-        try self.emit(.jalr(.zero, .ra, 0));
+        try self.emit(.addi(.a0, val, 0));
+        try self.epilogue();
     }
 
     fn genAssignLocal(self: *Self, assign_local: compile.Instr.AssignLocal) !void {
@@ -765,9 +749,26 @@ const ProcedureContext = struct {
         }
     }
 
+    fn prologue(self: *Self) !void {
+        try self.emit(.addi(.sp, .sp, -8));
+        try self.emit(.sd(.sp, 0, .ra));
+    }
+
+    fn epilogue(self: *Self) !void {
+        try self.emit(.ld(.ra, .sp, 0));
+        try self.emit(.addi(.sp, .sp, 8));
+        try self.emit(.jalr(.zero, .ra, 0));
+    }
+
     fn codegenProc(self: *Self, proc: compile.Procedure) !void {
+        var first = true;
         for (proc.blocks) |block| {
             try self.ctx.block_addrs.putNoClobber(block.ref, self.ctx.instructions.items.len);
+            if (first) {
+                try self.prologue();
+                first = false;
+            }
+
             try self.codegenBlock(block);
         }
     }
@@ -845,19 +846,18 @@ fn codegenReadInt(self: *Context) !void {
     try self.emit(.jalr(.zero, .ra, 0));
 }
 
+fn codegenExit(self: *Context) !void {
+    try self.emit(.addi(.a7, .zero, @intFromEnum(std.os.linux.syscalls.RiscV64.exit)));
+    try self.emit(.ecall());
+    // This will never be run, but makes binary ninja understand that this exits the function.
+    try self.emit(.jalr(.zero, .ra, 0));
+}
+
 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_addrs = .init(allocator),
-        .print_block = @enumFromInt(maxBlockRef + 1),
-        .read_int_block = @enumFromInt(maxBlockRef + 2),
     };
     defer ctx.deinit();
 
@@ -873,10 +873,12 @@ pub fn create_elf(allocator: Allocator, mod: compile.Module) ![]u8 {
         std.debug.assert(proc_ctx.register_allocator.allocated.count() == 0);
     }
 
-    try ctx.block_addrs.putNoClobber(ctx.print_block, ctx.instructions.items.len);
+    try ctx.block_addrs.putNoClobber(mod.print_block, ctx.instructions.items.len);
     try codegenPrint(&ctx);
-    try ctx.block_addrs.putNoClobber(ctx.read_int_block, ctx.instructions.items.len);
+    try ctx.block_addrs.putNoClobber(mod.read_int_block, ctx.instructions.items.len);
     try codegenReadInt(&ctx);
+    try ctx.block_addrs.putNoClobber(mod.exit_block, ctx.instructions.items.len);
+    try codegenExit(&ctx);
 
     // TODO: make this less sheiße
     for (ctx.relocations.items) |relocation| {