Implement "return"

5 files changed, 41 insertions, 4 deletions

diff --git a/examples/4.src b/examples/4.src
index 7e6662e..e1ebcdc 100644
--- a/examples/4.src
+++ b/examples/4.src
@@ -1,3 +1,3 @@
-let hello = "hello";
+let x = 2;
 
-print(hello + " world!");
+return x+1;
diff --git a/grammar.ebnf b/grammar.ebnf
index 1ac139e..5340ba4 100644
--- a/grammar.ebnf
+++ b/grammar.ebnf
@@ -1,9 +1,11 @@
 Program ::= Statement+
 
-Statement ::= (VariableStatement | PrintStatement) SEMICOLON
+Statement ::= (VariableStatement | PrintStatement | ReturnStatement) SEMICOLON
 
 VariableStatement ::= ("let" IDENTIFIER | IDENTIFIER) EQUALS Expression
 
 PrintStatement :== PRINT LPAREN Expression RPAREN
 
+ReturnStatement :== RETURN Expression
+
 Expression :== NUMBER | IDENTIFIER | Expression + Expression
diff --git a/src/evaluator.zig b/src/evaluator.zig
index 9e57dbd..fd0c8e5 100644
--- a/src/evaluator.zig
+++ b/src/evaluator.zig
@@ -8,6 +8,7 @@ const EvaluatorError = error{
 
 pub const Evaluator = struct {
     variables: std.StringHashMap(?i64),
+    return_value: i64,
 
     allocator: std.mem.Allocator,
 
@@ -16,6 +17,7 @@ pub const Evaluator = struct {
         evaluator.* = .{
             .variables = std.StringHashMap(?i64).init(allocator),
             .allocator = allocator,
+            .return_value = 0,
         };
         return evaluator;
     }
@@ -34,7 +36,7 @@ pub const Evaluator = struct {
             try self.evaluate_statement(statement);
         }
 
-        return 0;
+        return self.return_value;
     }
 
     fn evaluate_statement(self: *Evaluator, statement: *parser.Node) !void {
@@ -44,6 +46,7 @@ pub const Evaluator = struct {
         return switch (statement.STATEMENT.statement.*) {
             .VARIABLE_STATEMENT => |*variable_statement| try self.evaluate_variable_statement(@ptrCast(variable_statement)),
             .PRINT_STATEMENT => |*print_statement| try self.evaluate_print_statement(@ptrCast(print_statement)),
+            .RETURN_STATEMENT => |*return_statement| try self.evaluate_return_statement(@ptrCast(return_statement)),
             else => unreachable,
         };
     }
@@ -79,6 +82,16 @@ pub const Evaluator = struct {
         std.debug.print("PRINT: {d}\n", .{print_value});
     }
 
+    // TODO: This is a functionless implementation of return, we should not do this
+    fn evaluate_return_statement(self: *Evaluator, return_statement: *parser.Node) !void {
+        errdefer std.debug.print("Error evaluating return statement\n", .{});
+        std.debug.assert(return_statement.* == parser.Node.RETURN_STATEMENT);
+
+        const return_value = try self.get_expression_value(return_statement.RETURN_STATEMENT.expression);
+
+        self.return_value = return_value;
+    }
+
     fn get_expression_value(self: *Evaluator, node: *parser.Node) !i64 {
         errdefer std.debug.print("Error getting statement value\n", .{});
         std.debug.assert(node.* == parser.Node.EXPRESSION);
diff --git a/src/parser.zig b/src/parser.zig
index e8fcd81..5d3a957 100644
--- a/src/parser.zig
+++ b/src/parser.zig
@@ -11,6 +11,7 @@ const NodeType = enum {
     STATEMENT,
     VARIABLE_STATEMENT,
     PRINT_STATEMENT,
+    RETURN_STATEMENT,
     EXPRESSION,
 };
 
@@ -29,6 +30,9 @@ pub const Node = union(NodeType) {
     PRINT_STATEMENT: struct {
         expression: *Node,
     },
+    RETURN_STATEMENT: struct {
+        expression: *Node,
+    },
     EXPRESSION: union(enum) {
         NUMBER: struct {
             value: i64,
@@ -84,6 +88,7 @@ pub const Parser = struct {
 
         const statement = switch (token) {
             .PRINT => try self.parse_print_statement(),
+            .RETURN => try self.parse_return_statement(),
             else => try self.parse_variable_statement(),
         };
 
@@ -138,6 +143,20 @@ pub const Parser = struct {
         });
     }
 
+    // ReturnStatement :== RETURN Expression
+    fn parse_return_statement(self: *Parser) ParserError!*Node {
+        errdefer std.debug.print("Error parsing return statement\n", .{});
+        _ = try self.accept_token(tokenizer.TokenType.RETURN);
+
+        const expression = try self.parse_expression();
+
+        return self.create_node(.{
+            .RETURN_STATEMENT = .{
+                .expression = @constCast(expression),
+            },
+        });
+    }
+
     // Expression :== NUMBER | IDENTIFIER | Expression + Expression
     fn parse_expression(self: *Parser) ParserError!*Node {
         errdefer std.debug.print("Error parsing expression\n", .{});
diff --git a/src/tokenizer.zig b/src/tokenizer.zig
index 20f80cc..aab8aa2 100644
--- a/src/tokenizer.zig
+++ b/src/tokenizer.zig
@@ -8,6 +8,7 @@ pub const TokenType = enum {
     // Keywords
     LET,
     PRINT,
+    RETURN,
 
     // Identifiers
     IDENTIFIER,
@@ -28,6 +29,7 @@ pub const TokenType = enum {
 pub const Token = union(TokenType) {
     LET: void,
     PRINT: void,
+    RETURN: void,
     IDENTIFIER: []u8,
     NUMBER: i64,
     EQUALS: void,
@@ -64,6 +66,7 @@ pub const Tokenizer = struct {
 
         if (std.mem.eql(u8, string, "let")) return Token{ .LET = void{} };
         if (std.mem.eql(u8, string, "print")) return Token{ .PRINT = void{} };
+        if (std.mem.eql(u8, string, "return")) return Token{ .RETURN = void{} };
 
         if (std.fmt.parseInt(i32, string, 10) catch null) |i| return Token{ .NUMBER = i };