continue continuing procedure calls

5 files changed, 128 insertions, 78 deletions

diff --git a/fibonacci.hgn b/fibonacci.hgn
index c42241b..362e244 100644
--- a/fibonacci.hgn
+++ b/fibonacci.hgn
@@ -1,4 +1,5 @@
 main := proc() {
+    p(69)
     a := 0
     b := 1
     n := 10
@@ -10,4 +11,9 @@ main := proc() {
         n = n - 1
     }
     print(a)
+    exit(0)
+}
+
+p := proc() {
+    print(1)
 }
diff --git a/src/Lexer.zig b/src/Lexer.zig
index 87ced92..4a4d895 100644
--- a/src/Lexer.zig
+++ b/src/Lexer.zig
@@ -21,6 +21,7 @@ pub const Token = struct {
         right_angle,
         left_angle_equal,
         right_angle_equal,
+        comma,
 
         // Keywords
         @"if",
@@ -86,6 +87,7 @@ fn getNext(self: *Self) Token {
         '-' => self.create(.minus),
         '=' => self.create(.equal),
         ':' => self.create(.colon),
+        ',' => self.create(.comma),
         '<' => if (self.eatIfEqual('='))
             self.create(.left_angle_equal)
         else
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| {
diff --git a/src/compile.zig b/src/compile.zig
index d2226f1..c84b2cc 100644
--- a/src/compile.zig
+++ b/src/compile.zig
@@ -32,6 +32,9 @@ fn IdCounter(Id: type) type {
 
 pub const Module = struct {
     procedures: []Procedure,
+    print_block: BlockRef,
+    read_int_block: BlockRef,
+    exit_block: BlockRef,
 
     pub fn format(self: Module, comptime fmt: []const u8, options: std.fmt.FormatOptions, writer: anytype) !void {
         _ = .{ fmt, options };
@@ -56,8 +59,8 @@ pub const Procedure = struct {
     pub fn format(self: Procedure, comptime fmt: []const u8, options: std.fmt.FormatOptions, writer: anytype) !void {
         _ = .{ fmt, options };
         try writer.print("{s}:\n", .{self.name});
-        for (self.blocks, 0..) |block, i| {
-            try writer.print("  ${}{}", .{ i, block });
+        for (self.blocks) |block| {
+            try writer.print("{}", .{block});
         }
         try writer.writeAll("\n");
     }
@@ -94,7 +97,7 @@ 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(":\n", .{});
+        try writer.print("  ${}:\n", .{@intFromEnum(self.ref)});
         for (self.instrs.items) |instr| {
             try writer.print("    {}\n", .{instr});
         }
@@ -113,7 +116,7 @@ pub const Instr = struct {
         jump: Jump,
         assign_local: AssignLocal,
         get_local: GetLocal,
-        exit: Exit,
+        ret: Return,
     };
 
     pub const Constant = struct {
@@ -148,12 +151,7 @@ pub const Instr = struct {
     pub const ProcCall = struct {
         dest: VReg,
         arg: VReg,
-        proc: Proc,
-
-        const Proc = enum {
-            print,
-            read_int,
-        };
+        proc: BlockRef,
 
         pub fn sources(self: ProcCall) Sources {
             return Sources.fromSlice(&.{self.arg}) catch unreachable;
@@ -200,11 +198,13 @@ pub const Instr = struct {
         }
     };
 
-    pub const Exit = struct {
+    pub const Return = struct {
+        val: VReg,
+
         pub const may_end_block = {};
 
-        pub fn sources(_: Exit) Sources {
-            return Sources.init(0) catch unreachable;
+        pub fn sources(self: Return) Sources {
+            return Sources.fromSlice(&.{self.val}) catch unreachable;
         }
     };
 
@@ -217,7 +217,7 @@ pub const Instr = struct {
     pub fn dest(self: *const Instr) ?VReg {
         return switch (self.type) {
             inline .constant, .bin_op, .proc_call, .get_local => |s| s.dest,
-            .branch, .jump, .exit, .assign_local => null,
+            .branch, .jump, .ret, .assign_local => null,
         };
     }
 
@@ -241,10 +241,10 @@ pub const Instr = struct {
                 },
             ),
             .proc_call => |proc_call| try writer.print(
-                "%{} = {s} %{}",
+                "%{} = call ${}(%{})",
                 .{
                     @intFromEnum(proc_call.dest),
-                    @tagName(proc_call.proc),
+                    @intFromEnum(proc_call.proc),
                     @intFromEnum(proc_call.arg),
                 },
             ),
@@ -265,7 +265,7 @@ pub const Instr = struct {
                 try writer.print("%{} = @{}", .{ @intFromEnum(get_local.dest), @intFromEnum(get_local.local) });
             },
 
-            .exit => |_| try writer.print("exit", .{}),
+            .ret => |ret| try writer.print("return %{}", .{@intFromEnum(ret.val)}),
         }
     }
 };
@@ -284,49 +284,75 @@ const CompileError = error{
 };
 
 pub fn compile(allocator: Allocator, source: []const u8, file: parse.File) !Module {
-    const procs = try allocator.alloc(Procedure, file.decls.len);
+    var block_ctr: IdCounter(BlockRef) = .init;
+    var procedures: std.StringHashMapUnmanaged(BlockRef) = .empty;
+    const first_blocks: []BlockRef = try allocator.alloc(BlockRef, file.decls.len);
+    for (file.decls, first_blocks) |decl, *first| {
+        const id = block_ctr.get();
+        try procedures.put(allocator, decl.inner.ident.getIdent(source), id);
+        first.* = id;
+    }
+
+    const print_block = block_ctr.get();
+    try procedures.put(allocator, "print", print_block);
+    const read_int_block = block_ctr.get();
+    try procedures.put(allocator, "read_int", read_int_block);
+    const exit_block = block_ctr.get();
+    try procedures.put(allocator, "exit", exit_block);
+
     var ctx: Context = .{
         .allocator = allocator,
-        .block_ctx = .init,
+        .source = source,
+        .block_ctr = block_ctr,
+        .procedures = procedures,
     };
-    for (procs, file.decls) |*proc, decl| {
+    const procs = try allocator.alloc(Procedure, file.decls.len);
+    for (procs, file.decls, first_blocks) |*proc, decl, first_block| {
         proc.* = try compileProcedure(
             &ctx,
-            source,
             decl.inner.ident,
             decl.inner.value.type.proc,
+            first_block,
         );
     }
     return .{
         .procedures = procs,
+        .print_block = print_block,
+        .read_int_block = read_int_block,
+        .exit_block = exit_block,
     };
 }
 
 const Context = struct {
     allocator: Allocator,
-    block_ctx: IdCounter(BlockRef),
+    source: []const u8,
+    block_ctr: IdCounter(BlockRef),
+    procedures: std.StringHashMapUnmanaged(BlockRef),
 };
 
 fn compileProcedure(
     ctx: *Context,
-    source: []const u8,
     name: Location,
     proc: parse.Expr.Type.Proc,
+    first_block: BlockRef,
 ) !Procedure {
     var pctx: ProcedureContext = .{
         .ctx = ctx,
-        .source = source,
         .vreg_ctr = .init,
         .lvar_ctr = .init,
         .scope = .{ .locals = .empty, .parent = null },
-        .blocks = .empty,
-        .current_block = undefined, // immediately set by `switchToNewBlock`
+        .blocks = try .init(ctx.allocator, &.{first_block}, &.{.{ .ref = first_block }}),
+        .current_block = first_block,
     };
-    _ = try pctx.switchToNewBlock();
     try pctx.compileBlock(proc.body);
+    const proc_res = pctx.vreg_ctr.get();
     try pctx.addInstr(.{
         .loc = .{ .start = 0, .end = 0 },
-        .type = .{ .exit = .{} },
+        .type = .{ .constant = .{ .dest = proc_res, .value = 0 } },
+    });
+    try pctx.addInstr(.{
+        .loc = .{ .start = 0, .end = 0 },
+        .type = .{ .ret = .{ .val = proc_res } },
     });
     var blocks = try ctx.allocator.alloc(BasicBlock, pctx.blocks.count());
     var it = pctx.blocks.iterator();
@@ -335,12 +361,11 @@ fn compileProcedure(
         std.debug.assert(kv.key_ptr.* == kv.value_ptr.ref);
         blocks[i] = kv.value_ptr.*;
     }
-    return try .init(ctx.allocator, name.getIdent(source), blocks);
+    return try .init(ctx.allocator, name.getIdent(ctx.source), blocks);
 }
 
 const ProcedureContext = struct {
     ctx: *Context,
-    source: []const u8,
 
     vreg_ctr: IdCounter(VReg),
     lvar_ctr: IdCounter(LVar),
@@ -374,13 +399,13 @@ const ProcedureContext = struct {
             .assign_var => |assign_var| {
                 const local = if (assign_var.is_decl) blk: {
                     const local = self.lvar_ctr.get();
-                    const name = assign_var.ident.getIdent(self.source);
+                    const name = assign_var.ident.getIdent(self.ctx.source);
                     try self.scope.locals.put(self.ctx.allocator, name, local);
                     break :blk local;
                 } else blk: {
                     var scope: ?*Scope = &self.scope;
                     while (scope) |s| : (scope = s.parent) {
-                        if (s.locals.get(assign_var.ident.getIdent(self.source))) |local|
+                        if (s.locals.get(assign_var.ident.getIdent(self.ctx.source))) |local|
                             break :blk local;
                     } else return error.UnknownVariable;
                 };
@@ -431,7 +456,7 @@ const ProcedureContext = struct {
         switch (expr.type) {
             .integer_literal => try self.addInstr(.{
                 .loc = expr.loc,
-                .type = .{ .constant = .{ .dest = dest, .value = expr.loc.getInt(self.source) } },
+                .type = .{ .constant = .{ .dest = dest, .value = expr.loc.getInt(self.ctx.source) } },
             }),
             .bin_op => |binop| {
                 const lhs = try self.compileExpr(binop.lhs);
@@ -457,7 +482,8 @@ const ProcedureContext = struct {
             },
             .call => |call| {
                 if (call.proc.type != .identifier) return error.CanOnlyCallIdentifiers;
-                const proc = std.meta.stringToEnum(Instr.ProcCall.Proc, call.proc.loc.getIdent(self.source)) orelse return error.UnknownProcedure;
+                const proc = self.ctx.procedures.get(call.proc.loc.getIdent(self.ctx.source)) orelse
+                    return error.UnknownProcedure;
 
                 const arg = try self.compileExpr(call.arg);
                 try self.addInstr(.{
@@ -473,10 +499,11 @@ const ProcedureContext = struct {
                 var scope: ?*Scope = &self.scope;
                 const local: LVar = blk: {
                     while (scope) |s| : (scope = s.parent) {
-                        if (s.locals.get(expr.loc.getIdent(self.source))) |local| {
+                        if (s.locals.get(expr.loc.getIdent(self.ctx.source))) |local| {
                             break :blk local;
                         }
                     }
+                    std.log.debug("{s}", .{expr.loc.getIdent(self.ctx.source)});
                     return error.UnknownVariable;
                 };
                 try self.addInstr(.{
@@ -531,7 +558,7 @@ const ProcedureContext = struct {
     }
 
     fn switchToNewBlock(self: *ProcedureContext) !BlockRef {
-        const ref = self.ctx.block_ctx.get();
+        const ref = self.ctx.block_ctr.get();
         try self.blocks.putNoClobber(self.ctx.allocator, ref, .{ .ref = ref });
         self.current_block = ref;
         return ref;
diff --git a/src/parse.zig b/src/parse.zig
index bd0ca46..451b050 100644
--- a/src/parse.zig
+++ b/src/parse.zig
@@ -32,8 +32,9 @@ pub const File = struct {
     };
 
     fn format(self: File, writer: anytype, source: []const u8, indent: usize) !void {
-        for (self.decls) |decl| {
-            try writer.print("{s} {s} {}", .{
+        for (0.., self.decls) |i, decl| {
+            try writer.print("{s}{s} {s} {}\n", .{
+                if (i == 0) "" else "\n",
                 decl.inner.ident.getIdent(source),
                 ":=",
                 fmt(decl.inner.value, source, indent),
@@ -182,6 +183,7 @@ const ParseError = error{
     UnexpectedToken,
     InvalidAssignTarget,
     ExprStatementMustBeCall,
+    ExpectedCOmmaOrIdentifier,
 };
 
 pub fn file(allocator: Allocator, lexer: *Lexer) !File {
@@ -272,8 +274,19 @@ fn parseProc(allocator: Allocator, lexer: *Lexer) ParseError!*Expr {
     if (lexer.peek().type != .proc) return parseComparisons(allocator, lexer);
     const proc = try mustEat(lexer, .proc);
     _ = try mustEat(lexer, .left_paren);
-    // TODO: parameters
-    _ = try mustEat(lexer, .right_paren);
+    var params: std.ArrayList(Lexer.Location) = .init(allocator);
+    while (true) {
+        const tok = lexer.next();
+        switch (tok.type) {
+            .right_paren => break,
+            .identifier => {
+                try params.append(tok.loc);
+                if (lexer.peek().type == .right_paren) continue;
+                _ = try mustEat(lexer, .comma);
+            },
+            else => return error.ExpectedCOmmaOrIdentifier,
+        }
+    }
     const body = try parseBlock(allocator, lexer);
 
     return allocate(Expr, allocator, .{