Revert "Deprecate `defer`"

This reverts commit 7e3e245f6809946ea06ef1998bcabb7e0902fbd7.

5 files changed, 67 insertions, 2 deletions

diff --git a/src/compile/blocks.c b/src/compile/blocks.c
index 7d53d44e..1059fd34 100644
--- a/src/compile/blocks.c
+++ b/src/compile/blocks.c
@@ -16,6 +16,7 @@ Text_t compile_block_expression(env_t *env, ast_t *ast) {
     if (stmts && !stmts->next) return compile(env, stmts->ast);
 
     Text_t code = Text("({\n");
+    deferral_t *prev_deferred = env->deferred;
     env = fresh_scope(env);
     for (ast_list_t *stmt = stmts; stmt; stmt = stmt->next)
         prebind_statement(env, stmt->ast);
@@ -23,6 +24,11 @@ Text_t compile_block_expression(env_t *env, ast_t *ast) {
         if (stmt->next) {
             code = Texts(code, compile_statement(env, stmt->ast), "\n");
         } else {
+            // TODO: put defer after evaluating block expression
+            for (deferral_t *deferred = env->deferred; deferred && deferred != prev_deferred;
+                 deferred = deferred->next) {
+                code = Texts(code, compile_statement(deferred->defer_env, deferred->block));
+            }
             code = Texts(code, compile(env, stmt->ast), ";\n");
         }
         bind_statement(env, stmt->ast);
@@ -37,6 +43,7 @@ Text_t compile_inline_block(env_t *env, ast_t *ast) {
 
     Text_t code = EMPTY_TEXT;
     ast_list_t *stmts = Match(ast, Block)->statements;
+    deferral_t *prev_deferred = env->deferred;
     env = fresh_scope(env);
     for (ast_list_t *stmt = stmts; stmt; stmt = stmt->next)
         prebind_statement(env, stmt->ast);
@@ -44,5 +51,8 @@ Text_t compile_inline_block(env_t *env, ast_t *ast) {
         code = Texts(code, compile_statement(env, stmt->ast), "\n");
         bind_statement(env, stmt->ast);
     }
+    for (deferral_t *deferred = env->deferred; deferred && deferred != prev_deferred; deferred = deferred->next) {
+        code = Texts(code, compile_statement(deferred->defer_env, deferred->block));
+    }
     return code;
 }
diff --git a/src/compile/expressions.c b/src/compile/expressions.c
index ceae1b61..a4603cd5 100644
--- a/src/compile/expressions.c
+++ b/src/compile/expressions.c
@@ -231,6 +231,7 @@ Text_t compile(env_t *env, ast_t *ast) {
         else return compile_statement(env, ast);
     }
     case Use: code_err(ast, "Compiling 'use' as expression!");
+    case Defer: code_err(ast, "Compiling 'defer' as expression!");
     case TableEntry: code_err(ast, "Table entries should not be compiled directly");
     case Declare:
     case Assign:
diff --git a/src/compile/functions.c b/src/compile/functions.c
index e3dbc2e7..d2320e04 100644
--- a/src/compile/functions.c
+++ b/src/compile/functions.c
@@ -252,6 +252,7 @@ Text_t compile_lambda(env_t *env, ast_t *ast) {
     Text_t name = namespace_name(env, env->namespace, Texts("lambda$", lambda->id));
 
     env_t *body_scope = fresh_scope(env);
+    body_scope->deferred = NULL;
     for (arg_ast_t *arg = lambda->args; arg; arg = arg->next) {
         type_t *arg_type = get_arg_ast_type(env, arg);
         set_binding(body_scope, arg->name, arg_type, Texts("_$", arg->name));
@@ -317,6 +318,9 @@ Text_t compile_lambda(env_t *env, ast_t *ast) {
         else body = Texts(body, compile_statement(body_scope, FakeAST(Return, stmt->ast)), "\n");
         bind_statement(body_scope, stmt->ast);
     }
+    if ((ret_t->tag == VoidType || ret_t->tag == AbortType) && body_scope->deferred)
+        body = Texts(body, compile_statement(body_scope, FakeAST(Return)), "\n");
+
     env->code->lambdas = Texts(env->code->lambdas, code, " {\n", body, "\n}\n");
     return Texts("((Closure_t){", name, ", ", userdata, "})");
 }
@@ -531,6 +535,10 @@ static void add_closed_vars(Table_t *closed_vars, env_t *enclosing_scope, env_t
         add_closed_vars(closed_vars, enclosing_scope, scope, reduction->key ? reduction->key : item);
         break;
     }
+    case Defer: {
+        add_closed_vars(closed_vars, enclosing_scope, env, Match(ast, Defer)->body);
+        break;
+    }
     case Return: {
         ast_t *ret = Match(ast, Return)->value;
         if (ret) add_closed_vars(closed_vars, enclosing_scope, env, ret);
@@ -656,6 +664,7 @@ Text_t compile_function(env_t *env, Text_t name_code, ast_t *ast, Text_t *static
         body_scope->namespace = body_scope->namespace->parent;
     }
 
+    body_scope->deferred = NULL;
     for (arg_ast_t *arg = args; arg; arg = arg->next) {
         type_t *arg_type = get_arg_ast_type(env, arg);
         set_binding(body_scope, arg->name, arg_type, Texts("_$", arg->name));
diff --git a/src/compile/loops.c b/src/compile/loops.c
index 921f879e..588be4f0 100644
--- a/src/compile/loops.c
+++ b/src/compile/loops.c
@@ -42,6 +42,7 @@ Text_t compile_for_loop(env_t *env, ast_t *ast) {
     loop_ctx_t loop_ctx = (loop_ctx_t){
         .loop_name = "for",
         .loop_vars = for_->vars,
+        .deferred = body_scope->deferred,
         .next = body_scope->loop_ctx,
     };
     body_scope->loop_ctx = &loop_ctx;
@@ -355,6 +356,7 @@ Text_t compile_repeat(env_t *env, ast_t *ast) {
     env_t *scope = fresh_scope(env);
     loop_ctx_t loop_ctx = (loop_ctx_t){
         .loop_name = "repeat",
+        .deferred = scope->deferred,
         .next = env->loop_ctx,
     };
     scope->loop_ctx = &loop_ctx;
@@ -371,6 +373,7 @@ Text_t compile_while(env_t *env, ast_t *ast) {
     env_t *scope = fresh_scope(env);
     loop_ctx_t loop_ctx = (loop_ctx_t){
         .loop_name = "while",
+        .deferred = scope->deferred,
         .next = env->loop_ctx,
     };
     scope->loop_ctx = &loop_ctx;
@@ -396,7 +399,11 @@ Text_t compile_skip(env_t *env, ast_t *ast) {
                 ctx->skip_label = Texts("skip_", skip_label_count);
                 ++skip_label_count;
             }
-            return Texts("goto ", ctx->skip_label, ";");
+            Text_t code = EMPTY_TEXT;
+            for (deferral_t *deferred = env->deferred; deferred && deferred != ctx->deferred; deferred = deferred->next)
+                code = Texts(code, compile_statement(deferred->defer_env, deferred->block));
+            if (code.length > 0) return Texts("{\n", code, "goto ", ctx->skip_label, ";\n}\n");
+            else return Texts("goto ", ctx->skip_label, ";");
         }
     }
     if (env->loop_ctx) code_err(ast, "This is not inside any loop");
@@ -418,7 +425,11 @@ Text_t compile_stop(env_t *env, ast_t *ast) {
                 ctx->stop_label = Texts("stop_", stop_label_count);
                 ++stop_label_count;
             }
-            return Texts("goto ", ctx->stop_label, ";");
+            Text_t code = EMPTY_TEXT;
+            for (deferral_t *deferred = env->deferred; deferred && deferred != ctx->deferred; deferred = deferred->next)
+                code = Texts(code, compile_statement(deferred->defer_env, deferred->block));
+            if (code.length > 0) return Texts("{\n", code, "goto ", ctx->stop_label, ";\n}\n");
+            else return Texts("goto ", ctx->stop_label, ";");
         }
     }
     if (env->loop_ctx) code_err(ast, "This is not inside any loop");
diff --git a/src/compile/statements.c b/src/compile/statements.c
index 2c680b3b..37eff680 100644
--- a/src/compile/statements.c
+++ b/src/compile/statements.c
@@ -84,11 +84,40 @@ static Text_t _compile_statement(env_t *env, ast_t *ast) {
     case Skip: return compile_skip(env, ast);
     case Stop: return compile_stop(env, ast);
     case Pass: return Text(";");
+    case Defer: {
+        ast_t *body = Match(ast, Defer)->body;
+        Table_t closed_vars = get_closed_vars(env, NULL, body);
+
+        static int defer_id = 0;
+        env_t *defer_env = fresh_scope(env);
+        Text_t code = EMPTY_TEXT;
+        for (int64_t i = 0; i < closed_vars.entries.length; i++) {
+            struct {
+                const char *name;
+                binding_t *b;
+            } *entry = closed_vars.entries.data + closed_vars.entries.stride * i;
+            if (entry->b->type->tag == ModuleType) continue;
+            if (Text$starts_with(entry->b->code, Text("userdata->"), NULL)) {
+                Table$str_set(defer_env->locals, entry->name, entry->b);
+            } else {
+                Text_t defer_name = Texts("defer$", ++defer_id, "$", entry->name);
+                defer_id += 1;
+                code = Texts(code, compile_declaration(entry->b->type, defer_name), " = ", entry->b->code, ";\n");
+                set_binding(defer_env, entry->name, entry->b->type, defer_name);
+            }
+        }
+        env->deferred = new (deferral_t, .defer_env = defer_env, .block = body, .next = env->deferred);
+        return code;
+    }
     case Return: {
         if (!env->fn) code_err(ast, "This return statement is not inside any function");
         ast_t *ret = Match(ast, Return)->value;
 
         Text_t code = EMPTY_TEXT;
+        for (deferral_t *deferred = env->deferred; deferred; deferred = deferred->next) {
+            code = Texts(code, compile_statement(deferred->defer_env, deferred->block));
+        }
+
         type_t *ret_type = get_function_return_type(env, env->fn);
         if (ret) {
             if (ret_type->tag == VoidType || ret_type->tag == AbortType)
@@ -104,6 +133,11 @@ static Text_t _compile_statement(env_t *env, ast_t *ast) {
                 }
             }
             Text_t value = compile_to_type(env, ret, ret_type);
+            if (env->deferred) {
+                code = Texts(compile_declaration(ret_type, Text("ret")), " = ", value, ";\n", code);
+                value = Text("ret");
+            }
+
             return Texts(code, "return ", value, ";");
         } else {
             if (ret_type->tag != VoidType)