Fix for #79

7 files changed, 224 insertions, 16 deletions

diff --git a/README.md b/README.md
index 86018f7..aa773f4 100644
--- a/README.md
+++ b/README.md
@@ -422,12 +422,14 @@ println "You are {age} years old.";
 Zen C allows manual memory management with ergonomic aids.
 
 #### Defer
-Execute code when the current scope exits.
+Execute code when the current scope exits. Defer statements are executed in LIFO (last-in, first-out) order.
 ```zc
 var f = fopen("file.txt", "r");
 defer fclose(f);
 ```
 
+> To prevent undefined behavior, control flow statements (`return`, `break`, `continue`, `goto`) are **not allowed** inside a `defer` block.
+
 #### Autofree
 Automatically free the variable when scope exits.
 ```zc
diff --git a/src/codegen/codegen.c b/src/codegen/codegen.c
index 776ed80..457594f 100644
--- a/src/codegen/codegen.c
+++ b/src/codegen/codegen.c
@@ -1920,12 +1920,18 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
         codegen_node_single(ctx, node->guard_stmt.body, out);
         break;
     case NODE_WHILE:
+    {
+        loop_defer_boundary[loop_depth++] = defer_count;
         fprintf(out, "while (");
         codegen_expression(ctx, node->while_stmt.condition, out);
         fprintf(out, ") ");
         codegen_node_single(ctx, node->while_stmt.body, out);
+        loop_depth--;
         break;
+    }
     case NODE_FOR:
+    {
+        loop_defer_boundary[loop_depth++] = defer_count;
         fprintf(out, "for (");
         if (node->for_stmt.init)
         {
@@ -1963,8 +1969,19 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
         }
         fprintf(out, ") ");
         codegen_node_single(ctx, node->for_stmt.body, out);
+        loop_depth--;
         break;
+    }
     case NODE_BREAK:
+        // Run defers from current scope down to loop boundary before breaking
+        if (loop_depth > 0)
+        {
+            int boundary = loop_defer_boundary[loop_depth - 1];
+            for (int i = defer_count - 1; i >= boundary; i--)
+            {
+                codegen_node_single(ctx, defer_stack[i], out);
+            }
+        }
         if (node->break_stmt.target_label)
         {
             fprintf(out, "goto __break_%s;\n", node->break_stmt.target_label);
@@ -1975,6 +1992,15 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
         }
         break;
     case NODE_CONTINUE:
+        // Run defers from current scope down to loop boundary before continuing
+        if (loop_depth > 0)
+        {
+            int boundary = loop_defer_boundary[loop_depth - 1];
+            for (int i = defer_count - 1; i >= boundary; i--)
+            {
+                codegen_node_single(ctx, defer_stack[i], out);
+            }
+        }
         if (node->continue_stmt.target_label)
         {
             fprintf(out, "goto __continue_%s;\n", node->continue_stmt.target_label);
@@ -2001,23 +2027,39 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
         fprintf(out, "%s:;\n", node->label_stmt.label_name);
         break;
     case NODE_DO_WHILE:
+    {
+        loop_defer_boundary[loop_depth++] = defer_count;
         fprintf(out, "do ");
         codegen_node_single(ctx, node->do_while_stmt.body, out);
         fprintf(out, " while (");
         codegen_expression(ctx, node->do_while_stmt.condition, out);
         fprintf(out, ");\n");
+        loop_depth--;
         break;
+    }
     // Loop constructs: loop, repeat, for-in
     case NODE_LOOP:
+    {
         // loop { ... } => while (1) { ... }
+        loop_defer_boundary[loop_depth++] = defer_count;
         fprintf(out, "while (1) ");
         codegen_node_single(ctx, node->loop_stmt.body, out);
+        loop_depth--;
         break;
+    }
     case NODE_REPEAT:
+    {
+        loop_defer_boundary[loop_depth++] = defer_count;
         fprintf(out, "for (int _rpt_i = 0; _rpt_i < (%s); _rpt_i++) ", node->repeat_stmt.count);
         codegen_node_single(ctx, node->repeat_stmt.body, out);
+        loop_depth--;
         break;
+    }
     case NODE_FOR_RANGE:
+    {
+        // Track loop entry for defer boundary
+        loop_defer_boundary[loop_depth++] = defer_count;
+
         fprintf(out, "for (");
         if (strstr(g_config.cc, "tcc"))
         {
@@ -2049,7 +2091,10 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
             fprintf(out, "++) ");
         }
         codegen_node_single(ctx, node->for_range.body, out);
+
+        loop_depth--;
         break;
+    }
     case NODE_ASM:
     {
         int is_extended = (node->asm_stmt.num_outputs > 0 || node->asm_stmt.num_inputs > 0 ||
@@ -2229,9 +2274,9 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
     }
     case NODE_RETURN:
     {
+        int has_defers = (defer_count > func_defer_boundary);
         int handled = 0;
-        // If returning a variable that implements Drop, we must zero it out
-        // to prevent the cleanup attribute from destroying the resource we just returned.
+
         if (node->ret.value && node->ret.value->type == NODE_EXPR_VAR)
         {
             char *tname = infer_type(ctx, node->ret.value);
@@ -2261,7 +2306,13 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
                     codegen_expression(ctx, node->ret.value, out);
                     fprintf(out, "; memset(&");
                     codegen_expression(ctx, node->ret.value, out);
-                    fprintf(out, ", 0, sizeof(_z_ret_mv)); _z_ret_mv; });\n");
+                    fprintf(out, ", 0, sizeof(_z_ret_mv)); ");
+                    // Run defers before returning
+                    for (int i = defer_count - 1; i >= func_defer_boundary; i--)
+                    {
+                        codegen_node_single(ctx, defer_stack[i], out);
+                    }
+                    fprintf(out, "_z_ret_mv; });\n");
                     handled = 1;
                 }
                 free(tname);
@@ -2270,9 +2321,36 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
 
         if (!handled)
         {
-            fprintf(out, "    return ");
-            codegen_expression(ctx, node->ret.value, out);
-            fprintf(out, ";\n");
+            if (has_defers && node->ret.value)
+            {
+                // Save return value, run defers, then return
+                fprintf(out, "    { ");
+                emit_auto_type(ctx, node->ret.value, node->token, out);
+                fprintf(out, " _z_ret = ");
+                codegen_expression(ctx, node->ret.value, out);
+                fprintf(out, "; ");
+                for (int i = defer_count - 1; i >= func_defer_boundary; i--)
+                {
+                    codegen_node_single(ctx, defer_stack[i], out);
+                }
+                fprintf(out, "return _z_ret; }\n");
+            }
+            else if (has_defers)
+            {
+                // No return value, just run defers
+                for (int i = defer_count - 1; i >= func_defer_boundary; i--)
+                {
+                    codegen_node_single(ctx, defer_stack[i], out);
+                }
+                fprintf(out, "    return;\n");
+            }
+            else
+            {
+                // No defers, simple return
+                fprintf(out, "    return ");
+                codegen_expression(ctx, node->ret.value, out);
+                fprintf(out, ";\n");
+            }
         }
         break;
     }
diff --git a/src/codegen/codegen.h b/src/codegen/codegen.h
index f8fe318..caf7c0c 100644
--- a/src/codegen/codegen.h
+++ b/src/codegen/codegen.h
@@ -48,6 +48,11 @@ extern int defer_count;
 extern ASTNode *defer_stack[];
 extern ASTNode *g_current_lambda;
 
+// Defer boundary tracking for proper defer execution on break/continue/return
 #define MAX_DEFER 1024
+#define MAX_LOOP_DEPTH 64
+extern int loop_defer_boundary[]; // defer_count at each loop entry
+extern int loop_depth;            // current loop nesting depth
+extern int func_defer_boundary;   // defer_count at function entry
 
 #endif
diff --git a/src/codegen/codegen_utils.c b/src/codegen/codegen_utils.c
index af1c862..f712ca4 100644
--- a/src/codegen/codegen_utils.c
+++ b/src/codegen/codegen_utils.c
@@ -16,6 +16,10 @@ ASTNode *defer_stack[MAX_DEFER];
 int defer_count = 0;
 ASTNode *g_current_lambda = NULL;
 
+int loop_defer_boundary[MAX_LOOP_DEPTH];
+int loop_depth = 0;
+int func_defer_boundary = 0;
+
 // Helper to emit variable declarations with array types.
 void emit_var_decl_type(ParserContext *ctx, FILE *out, const char *type_str, const char *var_name)
 {
diff --git a/src/parser/parser.h b/src/parser/parser.h
index 5757fbc..f566979 100644
--- a/src/parser/parser.h
+++ b/src/parser/parser.h
@@ -263,10 +263,11 @@ struct ParserContext
     int extern_symbol_count;
 
     // Codegen state:
-    FILE *hoist_out;   // For plugins to hoist code to file scope
-    int skip_preamble; // If 1, codegen_node(NODE_ROOT) won't emit preamble
-    int is_repl;       // REPL mode flag
-    int has_async;     // Track if async features are used
+    FILE *hoist_out;    // For plugins to hoist code to file scope
+    int skip_preamble;  // If 1, codegen_node(NODE_ROOT) won't emit preamble
+    int is_repl;        // REPL mode flag
+    int has_async;      // Track if async features are used
+    int in_defer_block; // Track if currently parsing inside a defer block
 };
 
 // Token helpers
diff --git a/src/parser/parser_stmt.c b/src/parser/parser_stmt.c
index 4577405..6e55d08 100644
--- a/src/parser/parser_stmt.c
+++ b/src/parser/parser_stmt.c
@@ -502,7 +502,12 @@ ASTNode *parse_guard(ParserContext *ctx, Lexer *l)
 
 ASTNode *parse_defer(ParserContext *ctx, Lexer *l)
 {
-    lexer_next(l); // defer
+    Token defer_token = lexer_next(l); // defer
+
+    // Track that we're parsing inside a defer block
+    int prev_in_defer = ctx->in_defer_block;
+    ctx->in_defer_block = 1;
+
     ASTNode *s;
     if (lexer_peek(l).type == TOK_LBRACE)
     {
@@ -513,7 +518,11 @@ ASTNode *parse_defer(ParserContext *ctx, Lexer *l)
         s = ast_create(NODE_RAW_STMT);
         s->raw_stmt.content = consume_and_rewrite(ctx, l);
     }
+
+    ctx->in_defer_block = prev_in_defer;
+
     ASTNode *n = ast_create(NODE_DEFER);
+    n->token = defer_token;
     n->defer_stmt.stmt = s;
     return n;
 }
@@ -1465,7 +1474,6 @@ ASTNode *parse_type_alias(ParserContext *ctx, Lexer *l)
     lexer_next(l); // consume '='
 
     char *o = parse_type(ctx, l);
-    // printf("DEBUG: parse_type returned '%s'\n", o);
 
     lexer_next(l); // consume ';' (parse_type doesn't consume it? parse_type calls parse_type_formal
                    // which doesn't consume ;?)
@@ -1485,8 +1493,16 @@ ASTNode *parse_type_alias(ParserContext *ctx, Lexer *l)
 
 ASTNode *parse_return(ParserContext *ctx, Lexer *l)
 {
-    lexer_next(l); // eat 'return'
+    Token return_token = lexer_next(l); // eat 'return'
+
+    // Error if return is used inside a defer block
+    if (ctx->in_defer_block)
+    {
+        zpanic_at(return_token, "'return' is not allowed inside a 'defer' block");
+    }
+
     ASTNode *n = ast_create(NODE_RETURN);
+    n->token = return_token;
 
     int handled = 0;
 
@@ -2696,8 +2712,16 @@ ASTNode *parse_statement(ParserContext *ctx, Lexer *l)
         // Break with optional label: break; or break 'outer;
         if (strncmp(tk.start, "break", 5) == 0 && tk.len == 5)
         {
-            lexer_next(l);
+            Token break_token = lexer_next(l);
+
+            // Error if break is used inside a defer block
+            if (ctx->in_defer_block)
+            {
+                zpanic_at(break_token, "'break' is not allowed inside a 'defer' block");
+            }
+
             ASTNode *n = ast_create(NODE_BREAK);
+            n->token = break_token;
             n->break_stmt.target_label = NULL;
             // Check for 'label
             if (lexer_peek(l).type == TOK_CHAR)
@@ -2719,8 +2743,16 @@ ASTNode *parse_statement(ParserContext *ctx, Lexer *l)
         // Continue with optional label
         if (strncmp(tk.start, "continue", 8) == 0 && tk.len == 8)
         {
-            lexer_next(l);
+            Token continue_token = lexer_next(l);
+
+            // Error if continue is used inside a defer block
+            if (ctx->in_defer_block)
+            {
+                zpanic_at(continue_token, "'continue' is not allowed inside a 'defer' block");
+            }
+
             ASTNode *n = ast_create(NODE_CONTINUE);
+            n->token = continue_token;
             n->continue_stmt.target_label = NULL;
             if (lexer_peek(l).type == TOK_CHAR)
             {
@@ -2902,6 +2934,13 @@ ASTNode *parse_statement(ParserContext *ctx, Lexer *l)
         if (strncmp(tk.start, "goto", 4) == 0 && tk.len == 4)
         {
             Token goto_tok = lexer_next(l); // eat 'goto'
+
+            // Error if goto is used inside a defer block
+            if (ctx->in_defer_block)
+            {
+                zpanic_at(goto_tok, "'goto' is not allowed inside a 'defer' block");
+            }
+
             Token next = lexer_peek(l);
 
             // Computed goto: goto *ptr;
diff --git a/tests/features/test_defer_control_flow.zc b/tests/features/test_defer_control_flow.zc
new file mode 100644
index 0000000..40c7f37
--- /dev/null
+++ b/tests/features/test_defer_control_flow.zc
@@ -0,0 +1,79 @@
+
+test "defer_runs_on_break" {
+    var result = 0;
+    
+    for i in 0..5 {
+        defer { result = result + 1; }
+        
+        if i == 2 {
+            break;
+        }
+    }
+    
+    assert(result == 3, "defer should run on break");
+}
+
+test "defer_runs_on_continue" {
+    var result = 0;
+    
+    for i in 0..5 {
+        defer { result = result + 1; }
+        
+        if i == 2 {
+            continue;
+        }
+    }
+    
+    assert(result == 5, "defer should run on continue");
+}
+
+fn early_return(x: int) -> int {
+    var side_effect = 0;
+    defer { side_effect = 42; }
+    
+    if x > 0 {
+        return x + side_effect;
+    }
+    
+    return 0;
+}
+
+test "defer_runs_on_return" {
+    var result = early_return(10);
+    assert(result == 10, "early return should work with defer");
+}
+
+test "defer_lifo_order" {
+    var result = 0;
+    
+    for i in 0..2 {
+        defer { result = result * 10 + 1; }
+        defer { result = result * 10 + 2; }
+        
+        if i == 0 {
+            break;
+        }
+    }
+    
+    assert(result == 21, "defers should run in LIFO order");
+}
+
+test "nested_loops_defer" {
+    var outer_count = 0;
+    var inner_count = 0;
+    
+    for i in 0..3 {
+        defer { outer_count = outer_count + 1; }
+        
+        for j in 0..4 {
+            defer { inner_count = inner_count + 1; }
+            
+            if j == 2 {
+                break;
+            }
+        }
+    }
+    
+    assert(inner_count == 9, "inner loop defer should run correctly");
+    assert(outer_count == 3, "outer loop defer should run correctly");
+}