Further fix for #121

5 files changed, 73 insertions, 8 deletions

diff --git a/src/codegen/codegen_decl.c b/src/codegen/codegen_decl.c
index b82e1af..d525963 100644
--- a/src/codegen/codegen_decl.c
+++ b/src/codegen/codegen_decl.c
@@ -991,7 +991,14 @@ int emit_tests_and_runner(ParserContext *ctx, ASTNode *node, FILE *out)
         if (cur->type == NODE_TEST)
         {
             fprintf(out, "static void _z_test_%d() {\n", test_count);
+            int saved = defer_count;
             codegen_walker(ctx, cur->test_stmt.body, out);
+            // Run defers
+            for (int i = defer_count - 1; i >= saved; i--)
+            {
+                codegen_node_single(ctx, defer_stack[i], out);
+            }
+            defer_count = saved;
             fprintf(out, "}\n");
             test_count++;
         }
diff --git a/src/codegen/codegen_stmt.c b/src/codegen/codegen_stmt.c
index 542a3c0..003ce42 100644
--- a/src/codegen/codegen_stmt.c
+++ b/src/codegen/codegen_stmt.c
@@ -886,9 +886,27 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
                 }
 
                 ASTNode *def = find_struct_def(ctx, clean_type);
-                if (def && def->type_info && def->type_info->traits.has_drop)
+                int has_drop = (def && def->type_info && def->type_info->traits.has_drop);
+
+                if (has_drop)
                 {
-                    fprintf(out, "__attribute__((cleanup(%s__Drop_glue))) ", clean_type);
+                    // Drop Flag: int __z_drop_flag_name = 1;
+                    fprintf(out, "int __z_drop_flag_%s = 1; ", node->var_decl.name);
+
+                    // Synthesize Defer: if (__z_drop_flag_name) Name__Drop_drop(&name);
+                    ASTNode *defer_node = xmalloc(sizeof(ASTNode));
+                    defer_node->type = NODE_RAW_STMT;
+                    char *stmt_str =
+                        xmalloc(256 + strlen(node->var_decl.name) * 2 + strlen(clean_type));
+                    sprintf(stmt_str, "if (__z_drop_flag_%s) %s__Drop_glue(&%s);",
+                            node->var_decl.name, clean_type, node->var_decl.name);
+                    defer_node->raw_stmt.content = stmt_str;
+                    defer_node->line = node->line;
+
+                    if (defer_count < MAX_DEFER)
+                    {
+                        defer_stack[defer_count++] = defer_node;
+                    }
                 }
 
                 // Emit Variable with Type
@@ -930,9 +948,29 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
                     }
 
                     ASTNode *def = find_struct_def(ctx, clean_type);
-                    if (def && def->type_info && def->type_info->traits.has_drop)
+                    int has_drop = (def && def->type_info && def->type_info->traits.has_drop);
+
+                    if (has_drop)
                     {
-                        fprintf(out, "__attribute__((cleanup(%s__Drop_glue))) ", clean_type);
+                        // Drop Flag: int __z_drop_flag_name = 1;
+                        fprintf(out, "int __z_drop_flag_%s = 1; ", node->var_decl.name);
+
+                        // Synthesize Defer: if (__z_drop_flag_name) Name__Drop_drop(&name);
+                        ASTNode *defer_node = xmalloc(sizeof(ASTNode));
+                        defer_node->type = NODE_RAW_STMT;
+                        // Build string
+                        char *stmt_str =
+                            xmalloc(256 + strlen(node->var_decl.name) * 2 + strlen(clean_type));
+                        sprintf(stmt_str, "if (__z_drop_flag_%s) %s__Drop_glue(&%s);",
+                                node->var_decl.name, clean_type, node->var_decl.name);
+                        defer_node->raw_stmt.content = stmt_str;
+                        defer_node->line = node->line;
+
+                        // Push to defer stack
+                        if (defer_count < MAX_DEFER)
+                        {
+                            defer_stack[defer_count++] = defer_node;
+                        }
                     }
 
                     emit_var_decl_type(ctx, out, inferred, node->var_decl.name);
@@ -940,6 +978,7 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
                     fprintf(out, " = ");
                     codegen_expression(ctx, node->var_decl.init_expr, out);
                     fprintf(out, ";\n");
+
                     if (node->var_decl.init_expr &&
                         emit_move_invalidation(ctx, node->var_decl.init_expr, out))
                     {
@@ -1474,6 +1513,7 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
     }
     case NODE_REPL_PRINT:
     {
+        // Safe block for printing
         fprintf(out, "{ ");
         emit_auto_type(ctx, node->repl_print.expr, node->token, out);
         fprintf(out, " _zval = (");
@@ -1652,6 +1692,9 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
         fprintf(out, ");\n");
         break;
     }
+    case NODE_RAW_STMT:
+        fprintf(out, "    %s\n", node->raw_stmt.content);
+        break;
     default:
         codegen_expression(ctx, node, out);
         fprintf(out, ";\n");
diff --git a/src/codegen/codegen_utils.c b/src/codegen/codegen_utils.c
index 38ae409..cdb2110 100644
--- a/src/codegen/codegen_utils.c
+++ b/src/codegen/codegen_utils.c
@@ -738,6 +738,7 @@ int emit_move_invalidation(ParserContext *ctx, ASTNode *node, FILE *out)
     {
         if (node->type == NODE_EXPR_VAR)
         {
+            fprintf(out, "__z_drop_flag_%s = 0; ", node->var_ref.name);
             fprintf(out, "memset(&%s, 0, sizeof(%s))", node->var_ref.name, node->var_ref.name);
             return 1;
         }
diff --git a/src/parser/parser_expr.c b/src/parser/parser_expr.c
index 50d96f0..0951475 100644
--- a/src/parser/parser_expr.c
+++ b/src/parser/parser_expr.c
@@ -3599,7 +3599,10 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
                 free(print_code);
 
                 ASTNode *n = ast_create(NODE_RAW_STMT);
-                n->raw_stmt.content = final_code;
+                char *stmt_code = xmalloc(strlen(final_code) + 2);
+                sprintf(stmt_code, "%s;", final_code);
+                free(final_code);
+                n->raw_stmt.content = stmt_code;
                 return n;
             }
         }
@@ -3639,7 +3642,10 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
             free(inner);
 
             ASTNode *n = ast_create(NODE_RAW_STMT);
-            n->raw_stmt.content = code;
+            char *stmt_code = xmalloc(strlen(code) + 2);
+            sprintf(stmt_code, "%s;", code);
+            free(code);
+            n->raw_stmt.content = stmt_code;
             return n;
         }
     }
diff --git a/src/parser/parser_stmt.c b/src/parser/parser_stmt.c
index f7c1d32..1b3a5d9 100644
--- a/src/parser/parser_stmt.c
+++ b/src/parser/parser_stmt.c
@@ -1887,7 +1887,11 @@ ASTNode *parse_statement(ParserContext *ctx, Lexer *l)
             }
 
             ASTNode *n = ast_create(NODE_RAW_STMT);
-            n->raw_stmt.content = code;
+            // Append semicolon to Statement Expression to make it a valid statement
+            char *stmt_code = xmalloc(strlen(code) + 2);
+            sprintf(stmt_code, "%s;", code);
+            free(code);
+            n->raw_stmt.content = stmt_code;
             n->raw_stmt.used_symbols = used_syms;
             n->raw_stmt.used_symbol_count = used_count;
             free(inner);
@@ -2433,7 +2437,11 @@ ASTNode *parse_statement(ParserContext *ctx, Lexer *l)
             }
 
             ASTNode *n = ast_create(NODE_RAW_STMT);
-            n->raw_stmt.content = code;
+            // Append semicolon to Statement Expression to make it a valid statement
+            char *stmt_code = xmalloc(strlen(code) + 2);
+            sprintf(stmt_code, "%s;", code);
+            free(code);
+            n->raw_stmt.content = stmt_code;
             n->raw_stmt.used_symbols = used_syms;
             n->raw_stmt.used_symbol_count = used_count;
             return n;