Allow empty return statements for void functions

2 files changed, 15 insertions, 4 deletions

diff --git a/src/generator.c b/src/generator.c
index 0b3bc4b..59a193e 100644
--- a/src/generator.c
+++ b/src/generator.c
@@ -333,7 +333,10 @@ void generate_block(Node *block, FILE *out);
 void generate_statement(Node *stmt, FILE *out)
 {
     if (stmt->type == AST_RETURN) {
-        generate_expr_into_rax(stmt->unary_expr, out);
+        if (stmt->unary_expr)
+            generate_expr_into_rax(stmt->unary_expr, out);
+        else
+            fprintf(out, "    xor rax, rax\n"); // Return 0
         fprintf(out, "    push rax\n"); // Save the return value
 
         // Run all the defer statements
diff --git a/src/parser.c b/src/parser.c
index 0294dc2..ad5f4ea 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -727,15 +727,23 @@ Node *parse_statement(Lexer *lexer)
     if (token.type == TOKEN_RETURN) {
         assert_token(Lexer_next(lexer), TOKEN_RETURN);
         node = Node_new(AST_RETURN);
-        node->unary_expr = parse_expression(lexer);
+        
+        token = Lexer_peek(lexer);
+        if (token.type == TOKEN_SEMICOLON) {
+            Lexer_next(lexer);
+            if (current_function->func.return_type->type != TYPE_VOID)
+                die_location(token.loc, "Return statement with no value in function that expects it");
+            return node;
+        } else {
+            node->unary_expr = parse_expression(lexer);
+            assert_token(Lexer_next(lexer), TOKEN_SEMICOLON);
+        }
 
         if (!is_convertible(node->unary_expr->expr_type, current_function->func.return_type)) {
             fprintf(stderr, "- Expected type: %s\n", type_to_str(node->unary_expr->expr_type));
             fprintf(stderr, "- Actual: %s\n", type_to_str(current_function->func.return_type));
             die_location(token.loc, "Return expression does not match function's return type");
         }
-
-        assert_token(Lexer_next(lexer), TOKEN_SEMICOLON);
     } else if (token.type == TOKEN_IF) {
         Lexer_next(lexer);
         node = Node_new(AST_IF);