From 7c853f83c9ee0180a178767d549a3df446c63b8d Mon Sep 17 00:00:00 2001
From: Bruce Hill <bruce@bruce-hill.com>
Date: Wed, 9 Oct 2024 00:45:52 -0400
Subject: Fix closure handling behavior for `defer` so that it can mutate
 closed-over variables.

---
 compile.c | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

(limited to 'compile.c')

diff --git a/compile.c b/compile.c
index 034a4772..ab596e82 100644
--- a/compile.c
+++ b/compile.c
@@ -947,10 +947,14 @@ CORD compile_statement(env_t *env, ast_t *ast)
             struct { const char *name; binding_t *b; } *entry = Table$entry(*closed_vars, i);
             if (entry->b->type->tag == ModuleType)
                 continue;
-            CORD defer_name = CORD_asprintf("defer$%d$%s", ++defer_id, entry->name);
-            code = CORD_all(
-                code, compile_declaration(entry->b->type, defer_name), " = ", entry->b->code, ";\n");
-            set_binding(defer_env, entry->name, new(binding_t, .type=entry->b->type, .code=defer_name));
+            if (CORD_ncmp(entry->b->code, 0, "userdata->", 0, strlen("userdata->")) == 0) {
+                set_binding(defer_env, entry->name, entry->b);
+            } else {
+                CORD defer_name = CORD_asprintf("defer$%d$%s", ++defer_id, entry->name);
+                code = CORD_all(
+                    code, compile_declaration(entry->b->type, defer_name), " = ", entry->b->code, ";\n");
+                set_binding(defer_env, entry->name, new(binding_t, .type=entry->b->type, .code=defer_name));
+            }
 
             if (env->fn_ctx->closed_vars)
                 Table$str_set(env->fn_ctx->closed_vars, entry->name, entry->b);
@@ -1001,6 +1005,12 @@ CORD compile_statement(env_t *env, ast_t *ast)
                     code_err(ast, "This function expects a return value of type %T, but this return has type %T", 
                              env->fn_ctx->return_type, ret_value_t);
             }
+
+            if (env->deferred) {
+                code = CORD_all(compile_declaration(env->fn_ctx->return_type, "ret"), " = ", value, ";\n", code);
+                value = "ret";
+            }
+
             return CORD_all(code, "return ", value, ";");
         } else {
             if (env->fn_ctx->return_type->tag != VoidType)
-- 
cgit v1.2.3