1 files changed, 54 insertions, 52 deletions

diff --git a/parser.c b/parser.c
index 590ed2b..7de4ce9 100644
--- a/parser.c
+++ b/parser.c
@@ -48,6 +48,42 @@ static void parse_type(struct type_node* p_node) {
     p_node->type = tok.data.ident;
 }
 
+static void parse_return(struct return_node* p_node) {
+    expect(TK_IDENT);
+    if (strcmp(tok.data.ident, "return") != 0)
+        PARSER_PANIC("unexpected token %s; expected: return", tok.data.ident);
+
+    if (!lexer_peek(&tok))
+        PARSER_PANIC("unexpected EOF in return statement");
+
+    if (tok.type == TK_SEMI) {
+        p_node->ret_val = NULL;
+        return;
+    }
+
+    p_node->ret_val = protected_alloc(sizeof(struct expr_node));
+    handle_expr(p_node->ret_val);
+}
+
+static void parse_int_lit(struct int_lit_node* p_node) {
+    expect(TK_INT_LIT);
+    p_node->val = tok.data.int_lit;
+}
+
+static void handle_expr(struct expr_node* p_node) {
+    switch (tok.type) {
+        case TK_SEMI:
+            p_node->type = EXPR_EMPTY;
+            return;
+        case TK_INT_LIT:
+            p_node->type = EXPR_INT_LIT;
+            parse_int_lit(&p_node->as._int_lit);
+            return;
+        default:
+            PARSER_PANIC("expected expression");
+    }
+}
+
 static void parse_var_decl(struct var_decl_node* p_node) {
     parse_type(&p_node->type);
     expect(TK_IDENT);
@@ -73,14 +109,25 @@ static void parse_group(struct group_node* p_node) {
 }
 
 static void handle_stmt(struct stmt_node* p_node) {
-    if (tok.type == TK_LCURLY) {
-        p_node->type = STMT_GROUP;
-        parse_group(&p_node->as._group);
-    } else {
-        p_node->type = STMT_EXPR;
-        handle_expr(&p_node->as._expr);
-        expect(TK_SEMI);
+    switch (tok.type) {
+        case TK_LCURLY:
+            p_node->type = STMT_GROUP;
+            parse_group(&p_node->as._group);
+            return;
+        case TK_IDENT:
+            if (strcmp(tok.data.ident, "return") == 0) {
+                p_node->type = STMT_RETURN;
+                parse_return(&p_node->as._return);
+            } else {
+                p_node->type = STMT_VAR_DECL;
+                parse_var_decl(&p_node->as._var_decl);
+            }
+            break;
+        default:
+            p_node->type = STMT_EXPR;
+            handle_expr(&p_node->as._expr);
     }
+    expect(TK_SEMI);
 }
 
 static void parse_arg_list(struct var_decl_node** pp_arg) {
@@ -116,51 +163,6 @@ static void parse_fn_decl(struct fn_decl_node* p_node) {
     parse_group(&p_node->body);
 }
 
-static void parse_return(struct return_node* p_node) {
-    expect(TK_IDENT);
-    if (strcmp(tok.data.ident, "return") != 0)
-        PARSER_PANIC("unexpected token %s; expected: return", tok.data.ident);
-
-    if (!lexer_peek(&tok))
-        PARSER_PANIC("unexpected EOF in return statement");
-
-    if (tok.type == TK_SEMI) {
-        p_node->ret_val = NULL;
-        return;
-    }
-
-    p_node->ret_val = protected_alloc(sizeof(struct expr_node));
-    handle_expr(p_node->ret_val);
-}
-
-static void parse_int_lit(struct int_lit_node* p_node) {
-    expect(TK_INT_LIT);
-    p_node->val = tok.data.int_lit;
-}
-
-static void handle_expr(struct expr_node* p_node) {
-    switch (tok.type) {
-        case TK_SEMI:
-            p_node->type = EXPR_EMPTY;
-            return;
-        case TK_IDENT:
-            if (strcmp(tok.data.ident, "return") == 0) {
-                p_node->type = EXPR_RETURN;
-                parse_return(&p_node->as._return);
-            } else {
-                p_node->type = EXPR_VAR_DECL;
-                parse_var_decl(&p_node->as._var_decl);
-            }
-            return;
-        case TK_INT_LIT:
-            p_node->type = EXPR_INT_LIT;
-            parse_int_lit(&p_node->as._int_lit);
-            return;
-        default:
-            PARSER_PANIC("expected expression");
-    }
-}
-
 static bool parse_root(struct root_node* p_node) {
     if (!lexer_peek(&tok)) return false;