1 files changed, 32 insertions, 44 deletions

diff --git a/src/parse.zig b/src/parse.zig
index 95a9bc9..4adeffb 100644
--- a/src/parse.zig
+++ b/src/parse.zig
@@ -44,20 +44,6 @@ pub const File = struct {
     }
 };
 
-pub const Block = struct {
-    loc: Location,
-    stmts: []Stmt,
-
-    fn format(self: Block, writer: anytype, source: []const u8, indent: usize) !void {
-        try writer.writeAll("{\n");
-        for (self.stmts) |stmt| {
-            try writer.print("{}\n", .{fmt(stmt, source, indent + 4)});
-        }
-        try writer.writeByteNTimes(' ', indent);
-        try writer.writeAll("}");
-    }
-};
-
 pub const Stmt = struct {
     loc: Location,
     type: Type,
@@ -65,7 +51,6 @@ pub const Stmt = struct {
     pub const Type = union(enum) {
         expr: *const Expr,
         assign_var: AssignVar,
-        block: Block,
         @"while": While,
 
         pub const AssignVar = struct {
@@ -76,7 +61,7 @@ pub const Stmt = struct {
 
         pub const While = struct {
             cond: *const Expr,
-            do: Block,
+            do: *const Expr,
         };
     };
 
@@ -84,7 +69,6 @@ pub const Stmt = struct {
         try writer.writeByteNTimes(' ', indent);
         return switch (self.type) {
             .expr => |expr| writer.print("{}", .{fmt(expr, source, indent)}),
-            .block => |b| writer.print("{}", .{fmt(b, source, indent)}),
             .assign_var => |assign_var| writer.print("{s} {s} {}", .{
                 assign_var.ident.getIdent(source),
                 if (assign_var.is_decl) ":=" else "=",
@@ -110,6 +94,7 @@ pub const Expr = struct {
         @"if": If,
         proc: Proc,
         @"return": Return,
+        block: Block,
 
         pub const BinOp = struct {
             lhs: *const Expr,
@@ -166,18 +151,22 @@ pub const Expr = struct {
 
         pub const If = struct {
             cond: *const Expr,
-            then: Block,
-            @"else": ?Block,
+            then: *const Expr,
+            @"else": ?*const Expr,
         };
 
         pub const Proc = struct {
-            body: Block,
+            body: *const Expr,
             param: ?Location,
         };
 
         pub const Return = struct {
             value: *const Expr,
         };
+
+        pub const Block = struct {
+            stmts: []Stmt,
+        };
     };
 
     fn format(self: Expr, writer: anytype, source: []const u8, indent: usize) !void {
@@ -201,6 +190,14 @@ pub const Expr = struct {
             },
             .proc => |proc| try writer.print("proc({s}) {}", .{ if (proc.param) |p| p.getIdent(source) else "", fmt(proc.body, source, indent) }),
             .@"return" => |ret| try writer.print("return {}", .{fmt(ret.value, source, indent)}),
+            .block => |block| {
+                try writer.writeAll("{\n");
+                for (block.stmts) |stmt| {
+                    try writer.print("{}\n", .{fmt(stmt, source, indent + 4)});
+                }
+                try writer.writeByteNTimes(' ', indent);
+                try writer.writeAll("}");
+            },
         }
     }
 };
@@ -230,32 +227,12 @@ pub fn file(allocator: Allocator, lexer: *Lexer) !File {
     };
 }
 
-fn parseBlock(allocator: Allocator, lexer: *Lexer) !Block {
-    const left_curly = try mustEat(lexer, .left_curly);
-    var stmts: std.ArrayList(Stmt) = .init(allocator);
-    while (lexer.peek().type != .right_curly) {
-        try stmts.append(try parseStatement(allocator, lexer));
-    }
-    const right_curly = try mustEat(lexer, .right_curly);
-    return .{
-        .loc = left_curly.loc.combine(right_curly.loc),
-        .stmts = try stmts.toOwnedSlice(),
-    };
-}
-
 fn parseStatement(allocator: Allocator, lexer: *Lexer) ParseError!Stmt {
     switch (lexer.peek().type) {
-        .left_curly => {
-            const b = try parseBlock(allocator, lexer);
-            return .{
-                .loc = b.loc,
-                .type = .{ .block = b },
-            };
-        },
         .@"while" => {
             const @"while" = lexer.next();
             const cond = try expression(allocator, lexer);
-            const do = try parseBlock(allocator, lexer);
+            const do = try expression(allocator, lexer);
             return .{
                 .loc = @"while".loc.combine(do.loc),
                 .type = .{ .@"while" = .{
@@ -318,7 +295,7 @@ fn parseProc(allocator: Allocator, lexer: *Lexer) ParseError!*Expr {
         else => return expected(&.{ .identifier, .right_paren }, param_tok),
     };
     if (param != null) _ = try mustEat(lexer, .right_paren);
-    const body = try parseBlock(allocator, lexer);
+    const body = try expression(allocator, lexer);
 
     return allocate(Expr, allocator, .{
         .loc = proc.loc.combine(body.loc),
@@ -354,10 +331,10 @@ fn parseIf(allocator: Allocator, lexer: *Lexer) !*Expr {
         .@"if" => {
             const @"if" = lexer.next();
             const cond = try expression(allocator, lexer);
-            const then = try parseBlock(allocator, lexer);
+            const then = try expression(allocator, lexer);
             const @"else" = if (lexer.peek().type == .@"else") blk: {
                 _ = lexer.next();
-                break :blk try parseBlock(allocator, lexer);
+                break :blk try expression(allocator, lexer);
             } else null;
             return try allocate(Expr, allocator, .{
                 .loc = @"if".loc.combine((@"else" orelse then).loc),
@@ -397,6 +374,17 @@ fn parsePrimaryExpr(allocator: Allocator, lexer: *Lexer) !*Expr {
                 return error.ExpectedRightParen;
             return res;
         },
+        .left_curly => {
+            var stmts: std.ArrayList(Stmt) = .init(allocator);
+            while (lexer.peek().type != .right_curly) {
+                try stmts.append(try parseStatement(allocator, lexer));
+            }
+            const right_curly = try mustEat(lexer, .right_curly);
+            return allocate(Expr, allocator, .{
+                .loc = token.loc.combine(right_curly.loc),
+                .type = .{ .block = .{ .stmts = try stmts.toOwnedSlice() } },
+            });
+        },
         .integer_literal => .{ .loc = token.loc, .type = .integer_literal },
         .identifier => .{ .loc = token.loc, .type = .identifier },
         else => return expected(&.{ .left_paren, .integer_literal, .identifier }, token),