2 files changed, 120 insertions, 10 deletions

diff --git a/boot/parse.c b/boot/parse.c
index 9f10d27..94b1e91 100644
--- a/boot/parse.c
+++ b/boot/parse.c
@@ -712,6 +712,65 @@ parser_statement_loop(struct Parser* p, struct Parser_Error* error)
 }
 
 struct Statement*
+parser_statement_block(struct Parser* p, struct Parser_Error* error)
+{
+    struct Block_Node block = CHECK(parser_block_node(p, error));
+    return statement_new(
+        STATEMENT_BLOCK, (union Statement_Value){ .block = { block } }, block.span, block.location);
+}
+
+struct Statement*
+parser_statement_return(struct Parser* p, struct Parser_Error* error)
+{
+    struct Token return_token = CHECK(parser_need(p, TOKEN_WORD_RETURN, error));
+
+    struct Expression* value = nil;
+    if (!token_ends_statement(&return_token)) { value = CHECK(parser_expression(p, error)); }
+
+    struct Span span = value ? return_token.span : span_merge(return_token.span, value->span);
+    union Statement_Value statement_value = { .return_value = { value } };
+    return statement_new(STATEMENT_RETURN, statement_value, span, return_token.location);
+}
+
+struct Statement*
+parser_statement_break(struct Parser* p, struct Parser_Error* error)
+{
+    struct Token break_token = CHECK(parser_need(p, TOKEN_WORD_BREAK, error));
+    struct Span span = break_token.span;
+    return statement_new(STATEMENT_BREAK, (union Statement_Value){ 0 }, span, break_token.location);
+}
+
+struct Statement*
+parser_statement_continue(struct Parser* p, struct Parser_Error* error)
+{
+    struct Token continue_token = CHECK(parser_need(p, TOKEN_WORD_CONTINUE, error));
+    struct Span span = continue_token.span;
+    return statement_new(
+        STATEMENT_CONTINUE, (union Statement_Value){ 0 }, span, continue_token.location);
+}
+
+struct Statement*
+parser_statement_defer(struct Parser* p, struct Parser_Error* error)
+{
+    struct Token defer_token = CHECK(parser_need(p, TOKEN_WORD_DEFER, error));
+
+    struct Span span = defer_token.span;
+    struct Statement_Value_Defer defer = { 0 };
+    if (parser_probe(p, TOKEN_CURLY_OPEN)) {
+        struct Block_Node block = CHECK(parser_block_node(p, error));
+        span = span_merge(span, block.span);
+        defer.block = block;
+    } else {
+        struct Expression* expression = CHECK(parser_expression(p, error));
+        span = span_merge(span, expression->span);
+        defer.expression = expression;
+    }
+
+    union Statement_Value value = { .defer = defer };
+    return statement_new(STATEMENT_DEFER, value, span, defer_token.location);
+}
+
+struct Statement*
 parser_statement(struct Parser* p, struct Parser_Error* error)
 {
     struct Token token = parser_peek(p);
@@ -724,16 +783,22 @@ parser_statement(struct Parser* p, struct Parser_Error* error)
 
     if (parser_could_be_variable_declaration(p)) return parser_statement_declaration(p, error);
 
-    if (token_is(&token, TOKEN_WORD_IF)) return parser_statement_conditional(p, error);
-
-    if (token_is(&token, TOKEN_WORD_FOR)) return parser_statement_loop(p, error);
-
-    if (token_is(&token, TOKEN_CURLY_OPEN)) {
-        // a block statement.
-        struct Block_Node block = CHECK(parser_block_node(p, error));
-        return statement_new(
-            STATEMENT_BLOCK, (union Statement_Value){ .block = { block } }, block.span,
-            block.location);
+    switch (token.kind) {
+    case TOKEN_WORD_IF:
+        return parser_statement_conditional(p, error);
+    case TOKEN_WORD_FOR:
+        return parser_statement_loop(p, error);
+    case TOKEN_CURLY_OPEN:
+        return parser_statement_block(p, error);
+    case TOKEN_WORD_RETURN:
+        return parser_statement_return(p, error);
+    case TOKEN_WORD_BREAK:
+        return parser_statement_break(p, error);
+    case TOKEN_WORD_CONTINUE:
+        return parser_statement_continue(p, error);
+    case TOKEN_WORD_DEFER:
+        return parser_statement_defer(p, error);
+    default: break;
     }
 
     struct Expression* expression = CHECK(parser_expression(p, error));
diff --git a/boot/tree.c b/boot/tree.c
index 2d49a2c..6e7fd6a 100644
--- a/boot/tree.c
+++ b/boot/tree.c
@@ -657,6 +657,10 @@ enum Statement_Kind
     STATEMENT_BLOCK,
     STATEMENT_CONDITIONAL,
     STATEMENT_LOOP,
+    STATEMENT_RETURN,
+    STATEMENT_BREAK,
+    STATEMENT_CONTINUE,
+    STATEMENT_DEFER,
 };
 
 struct Statement_Value_Expression
@@ -701,6 +705,20 @@ struct Statement_Value_Loop
     struct Block_Node body;
 };
 
+struct Statement_Value_Return
+{
+    // nil if there is no return value.
+    struct Expression* value;
+};
+
+struct Statement_Value_Defer
+{
+    // either a simple expression, or, if expression is nil,
+    // a block of code to execute.
+    struct Expression* expression;
+    struct Block_Node block;
+};
+
 union Statement_Value
 {
     struct Statement_Value_Expression expression;
@@ -708,6 +726,8 @@ union Statement_Value
     struct Statement_Value_Block block;
     struct Statement_Value_Conditional conditional;
     struct Statement_Value_Loop loop;
+    struct Statement_Value_Return return_value;
+    struct Statement_Value_Defer defer;
 };
 
 struct Statement
@@ -832,6 +852,31 @@ statement_print(const struct Statement* statement)
         block_node_print(&statement->value.loop.body);
         break;
     }
+    case STATEMENT_RETURN: {
+        printf("(return");
+        if (statement->value.return_value.value) {
+            printf(" ");
+            expression_print(statement->value.return_value.value);
+        }
+        printf(")");
+        break;
+    }
+    case STATEMENT_BREAK:
+        printf("(break)");
+        break;
+    case STATEMENT_CONTINUE:
+        printf("(continue)");
+        break;
+    case STATEMENT_DEFER: {
+        printf("(defer ");
+        if (statement->value.defer.expression) {
+            expression_print(statement->value.defer.expression);
+        } else {
+            block_node_print(&statement->value.defer.block);
+        }
+        printf(")");
+        break;
+    }
     default:
         failure("unexpected statement kind passed to `statement_print`");
         break;