5 files changed, 51 insertions, 8 deletions

diff --git a/recurse.hgn b/recurse.hgn
new file mode 100644
index 0000000..52b7da4
--- /dev/null
+++ b/recurse.hgn
@@ -0,0 +1,11 @@
+main := proc() {
+    print_up_to(10)
+    exit(0)
+}
+
+print_up_to := proc(n) {
+    print(n)
+    if n > 0 {
+        print_up_to(n - 1)
+    }
+}
diff --git a/src/Lexer.zig b/src/Lexer.zig
index 64457bf..9be99aa 100644
--- a/src/Lexer.zig
+++ b/src/Lexer.zig
@@ -29,6 +29,7 @@ pub const Token = struct {
         @"else",
         @"while",
         proc,
+        @"return",
     };
 };
 
@@ -144,7 +145,7 @@ fn identifierOrKeyword(self: *Self) Token {
         }
         const value = self.source[self.start..self.pos];
         return self.create(switch (std.meta.stringToEnum(Token.Type, value) orelse .invalid) {
-            .@"if", .@"else", .@"while", .proc => |t| t,
+            .@"if", .@"else", .@"while", .proc, .@"return" => |t| t,
             else => .identifier,
         });
     }
diff --git a/src/compile.zig b/src/compile.zig
index 8c734c7..0ce170c 100644
--- a/src/compile.zig
+++ b/src/compile.zig
@@ -597,6 +597,13 @@ const ProcedureContext = struct {
                 _ = proc;
                 return error.CannotDefineProcedureHere;
             },
+            .@"return" => |ret| {
+                const val = try self.compileExpr(ret.value);
+                try self.addInstr(.{
+                    .loc = expr.loc,
+                    .type = .{ .ret = .{ .val = val } },
+                });
+            },
         }
         return dest;
     }
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 {
diff --git a/triangle.hgn b/triangle.hgn
new file mode 100644
index 0000000..5459c12
--- /dev/null
+++ b/triangle.hgn
@@ -0,0 +1,10 @@
+main := proc() {
+    n := read_int(0)
+    print(triangle(n))
+    exit(0)
+}
+
+triangle := proc(n) {
+    if n == 0 { return 0 }
+    return triangle(n - 1) + n
+}