1 files changed, 21 insertions, 7 deletions
diff --git a/src/parse.zig b/src/parse.zig
index 3ecb93b..d8b4d15 100644
--- a/src/parse.zig
+++ b/src/parse.zig
@@ -109,6 +109,7 @@ pub const Expr = struct {
         identifier,
         @"if": If,
         proc: Proc,
+        @"return": Return,
 
         pub const BinOp = struct {
             lhs: *const Expr,
@@ -153,6 +154,10 @@ pub const Expr = struct {
             body: Block,
             param: ?Location,
         };
+
+        pub const Return = struct {
+            value: *const Expr,
+        };
     };
 
     fn format(self: Expr, writer: anytype, source: []const u8, indent: usize) !void {
@@ -174,9 +179,8 @@ pub const Expr = struct {
                     try writer.print(" else {}", .{fmt(@"else", source, indent)});
                 }
             },
-            .proc => |proc| {
-                try writer.print("proc() {}", .{fmt(proc.body, source, indent)});
-            },
+            .proc => |proc| try writer.print("proc() {}", .{fmt(proc.body, source, indent)}),
+            .@"return" => |ret| try writer.print("return {}", .{fmt(ret.value, source, indent)}),
         }
     }
 };
@@ -186,7 +190,7 @@ const ParseError = error{
     ExpectedRightParen,
     UnexpectedToken,
     InvalidAssignTarget,
-    ExprStatementMustBeCallOrIf,
+    InvalidExprStatement,
 };
 
 pub fn file(allocator: Allocator, lexer: *Lexer) !File {
@@ -258,8 +262,8 @@ fn parseStatement(allocator: Allocator, lexer: *Lexer) ParseError!Stmt {
                     .loc = lhs.loc.combine(value.loc),
                     .type = .{ .assign_var = .{ .ident = lhs.loc, .is_decl = colon != null, .value = value } },
                 };
-            } else if (lhs.type != .call and lhs.type != .@"if") {
-                return error.ExprStatementMustBeCallOrIf;
+            } else if (lhs.type != .call and lhs.type != .@"if" and lhs.type != .@"return") {
+                return error.InvalidExprStatement;
             }
             return .{
                 .loc = lhs.loc,
@@ -270,7 +274,17 @@ fn parseStatement(allocator: Allocator, lexer: *Lexer) ParseError!Stmt {
 }
 
 fn expression(allocator: Allocator, lexer: *Lexer) ParseError!*Expr {
-    return parseProc(allocator, lexer);
+    return parseReturn(allocator, lexer);
+}
+
+fn parseReturn(allocator: Allocator, lexer: *Lexer) ParseError!*Expr {
+    if (lexer.peek().type != .@"return") return parseProc(allocator, lexer);
+    const ret = try mustEat(lexer, .@"return");
+    const val = try parseProc(allocator, lexer);
+    return allocate(Expr, allocator, .{
+        .loc = ret.loc.combine(val.loc),
+        .type = .{ .@"return" = .{ .value = val } },
+    });
 }
 
 fn parseProc(allocator: Allocator, lexer: *Lexer) ParseError!*Expr {