improve error printing

2 files changed, 32 insertions, 18 deletions

diff --git a/src/main.zig b/src/main.zig
index d0e94db..78cc18c 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -6,7 +6,11 @@ pub const parse = @import("./parse.zig");
 pub const Lexer = @import("./Lexer.zig");
 pub const compile = @import("./compile.zig");
 
-pub fn main() !void {
+const blue = "\x1b[34m";
+const red = "\x1b[31m";
+const normal = "\x1b[0m";
+
+pub fn main() !u8 {
     var arena: std.heap.ArenaAllocator = .init(std.heap.smp_allocator);
     defer arena.deinit();
     const allocator = arena.allocator();
@@ -35,23 +39,33 @@ pub fn main() !void {
     defer allocator.free(source);
 
     var lexer: Lexer = .{ .source = source };
-    std.debug.print("Tokens:\n", .{});
+    std.debug.print(blue ++ "tokens:" ++ normal ++ "\n", .{});
     while (true) {
         const token = lexer.next();
         std.debug.print("    {}\n", .{token});
         if (token.type == .eof) break;
     }
     lexer = .{ .source = source };
-    const ast = try parse.file(allocator, &lexer);
-    std.debug.print("Parse tree:\n{}\n", .{parse.fmt(ast, source, 0)});
+    const ast = parse.file(allocator, &lexer) catch |err| {
+        std.debug.print(red ++ "parsing error: {}\n" ++ normal, .{err});
+        return 1;
+    };
+    std.debug.print(blue ++ "parse tree:" ++ normal ++ "\n{}\n", .{parse.fmt(ast, source, 0)});
     if (lexer.peek().type != .eof) {
-        std.debug.print("Unexpected token {}, expected end of file\n", .{lexer.next()});
+        std.debug.print(red ++ "Unexpected token {}, expected end of file\n" ++ normal, .{lexer.next()});
+        return 1;
     }
-    const module = try compile.compile(allocator, source, ast);
-    std.debug.print("Bytecode instructions:\n{}", .{module});
-    const elf = try codegen.create_elf(allocator, module);
+    const module = compile.compile(allocator, source, ast) catch |err| {
+        std.debug.print(red ++ "compilation error: {}\n" ++ normal, .{err});
+        return 1;
+    };
+    std.debug.print(blue ++ "bytecode instructions: " ++ normal ++ "\n{}", .{module});
+    const elf = codegen.create_elf(allocator, module) catch |err| {
+        std.debug.print(red ++ "code generation error: {}\n" ++ normal, .{err});
+        return 1;
+    };
     try out_file.writer().writeAll(elf);
-    std.debug.print("Run output:\n", .{});
+    std.debug.print(blue ++ "running program:" ++ normal ++ "\n", .{});
     if (run) {
         out_file.close();
 
@@ -62,7 +76,10 @@ pub fn main() !void {
             else
                 &.{ "qemu-riscv64", out_path.? },
         );
-        std.debug.print("{}\n{any}\n", .{ err, @errorReturnTrace() });
+        std.debug.print(red ++ "{}\n" ++ normal ++ "{any}\n", .{ err, @errorReturnTrace() });
+        return 1;
+    } else {
+        return 0;
     }
 }
 
diff --git a/src/parse.zig b/src/parse.zig
index 752eafb..f6fc0bc 100644
--- a/src/parse.zig
+++ b/src/parse.zig
@@ -295,7 +295,7 @@ fn parseProc(allocator: Allocator, lexer: *Lexer) ParseError!*Expr {
     const param: ?Location = switch (param_tok.type) {
         .identifier => param_tok.loc,
         .right_paren => null,
-        else => |ty| return expected(&.{ .identifier, .right_paren }, ty),
+        else => return expected(&.{ .identifier, .right_paren }, param_tok),
     };
     if (param != null) _ = try mustEat(lexer, .right_paren);
     const body = try parseBlock(allocator, lexer);
@@ -395,21 +395,18 @@ fn parsePrimaryExpr(allocator: Allocator, lexer: *Lexer) !*Expr {
         },
         .integer_literal => .{ .loc = token.loc, .type = .integer_literal },
         .identifier => .{ .loc = token.loc, .type = .identifier },
-        else => |t| {
-            std.debug.print("Expected '(', integer literal, or identifier. Got '{s}'\n", .{@tagName(t)});
-            return error.UnexpectedToken;
-        },
+        else => return expected(&.{ .left_paren, .integer_literal, .identifier }, token),
     });
 }
 
-fn expected(one_of: []const Token.Type, got: Token.Type) error{UnexpectedToken} {
+fn expected(one_of: []const Token.Type, got: Token) error{UnexpectedToken} {
     for (one_of) |ty|
-        std.debug.assert(ty != got);
+        std.debug.assert(ty != got.type);
     std.debug.print("Expected ", .{});
     for (0.., one_of) |i, ty| {
         std.debug.print("{s}{s}", .{ if (i == 0) "" else " or ", @tagName(ty) });
     }
-    std.debug.print(". Got {s}.\n", .{@tagName(got)});
+    std.debug.print(". Got {s} at {}.\n", .{ @tagName(got.type), got.loc });
     return error.UnexpectedToken;
 }