Revert "Deprecate `defer`"

This reverts commit 7e3e245f6809946ea06ef1998bcabb7e0902fbd7.

18 files changed, 187 insertions, 11 deletions

diff --git a/CHANGES.md b/CHANGES.md
index 95458289..c717a95e 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -20,7 +20,6 @@
   - Explicitly optional values can be declared as `my_var : T? = value`.
 - Deprecated `>> ... = ...` form of doctests. They are now called "debug logs"
   and you can specify multiple values: `>> a, b, c`
-- Deprecated `defer` statement
 - Struct fields that start with underscores can be accessed again and function
   arguments that start with underscore can be passed (but only as keyword
   arguments).
diff --git a/src/ast.c b/src/ast.c
index 7baf08b8..80127cf7 100644
--- a/src/ast.c
+++ b/src/ast.c
@@ -260,6 +260,7 @@ Text_t ast_to_sexp(ast_t *ast) {
         T(Skip, "(Skip ", quoted_text(data.target), ")");
         T(Stop, "(Stop ", quoted_text(data.target), ")");
         T(Pass, "(Pass)");
+        T(Defer, "(Defer ", ast_to_sexp(data.body), ")");
         T(Return, "(Return ", ast_to_sexp(data.value), ")");
         T(StructDef, "(StructDef \"", data.name, "\" ", arg_defs_to_sexp(data.fields), " ", ast_to_sexp(data.namespace),
           ")");
@@ -618,6 +619,10 @@ void ast_visit(ast_t *ast, void (*visitor)(ast_t *, void *), void *userdata) {
     case Skip:
     case Stop:
     case Pass: return;
+    case Defer: {
+        ast_visit(Match(ast, Defer)->body, visitor, userdata);
+        return;
+    }
     case Return: {
         ast_visit(Match(ast, Return)->value, visitor, userdata);
         return;
diff --git a/src/ast.h b/src/ast.h
index c1da167a..142bea19 100644
--- a/src/ast.h
+++ b/src/ast.h
@@ -262,6 +262,7 @@ typedef enum {
     Skip,
     Stop,
     Pass,
+    Defer,
     Return,
     StructDef,
     EnumDef,
@@ -407,6 +408,9 @@ struct ast_s {
         struct {
         } Pass;
         struct {
+            ast_t *body;
+        } Defer;
+        struct {
             ast_t *value;
         } Return;
         struct {
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)
diff --git a/src/environment.h b/src/environment.h
index ecb6fb4c..c726508d 100644
--- a/src/environment.h
+++ b/src/environment.h
@@ -14,10 +14,17 @@ typedef struct {
     Text_t variable_initializers;
 } compilation_unit_t;
 
+typedef struct deferral_s {
+    struct deferral_s *next;
+    struct env_s *defer_env;
+    ast_t *block;
+} deferral_t;
+
 typedef struct loop_ctx_s {
     struct loop_ctx_s *next;
     const char *loop_name;
     ast_list_t *loop_vars;
+    deferral_t *deferred;
     Text_t skip_label, stop_label;
 } loop_ctx_t;
 
@@ -38,6 +45,7 @@ typedef struct env_s {
     compilation_unit_t *code;
     ast_t *fn;
     loop_ctx_t *loop_ctx;
+    deferral_t *deferred;
     Closure_t *comprehension_action;
     bool do_source_mapping : 1;
     type_t *current_type;
diff --git a/src/formatter/formatter.c b/src/formatter/formatter.c
index 12f08810..2d220493 100644
--- a/src/formatter/formatter.c
+++ b/src/formatter/formatter.c
@@ -128,6 +128,8 @@ OptionalText_t format_inline_code(ast_t *ast, Table_t comments) {
         Text_t message = fmt_inline(assert->message, comments);
         return Texts("assert ", expr, ", ", message);
     }
+    /*inline*/ case Defer:
+        return Texts("defer ", fmt_inline(Match(ast, Defer)->body, comments));
     /*inline*/ case Lambda: {
         DeclareMatch(lambda, ast, Lambda);
         Text_t code = Texts("func(", format_inline_args(lambda->args, comments));
@@ -573,6 +575,8 @@ Text_t format_code(ast_t *ast, Table_t comments, Text_t indent) {
         DeclareMatch(extend, ast, Extend);
         return Texts("lang ", Text$from_str(extend->name), format_namespace(extend->body, comments, indent));
     }
+    /*multiline*/ case Defer:
+        return Texts("defer ", format_namespace(Match(ast, Defer)->body, comments, indent));
     /*multiline*/ case List: {
         if (inlined_fits) return inlined;
         ast_list_t *items = Match(ast, List)->items;
diff --git a/src/formatter/utils.c b/src/formatter/utils.c
index 0f638a57..9cd0227d 100644
--- a/src/formatter/utils.c
+++ b/src/formatter/utils.c
@@ -64,6 +64,7 @@ CONSTFUNC int suggested_blank_lines(ast_t *first, ast_t *second) {
     case While:
     case For:
     case Block:
+    case Defer:
     case ConvertDef:
     case FunctionDef:
     case Lambda:
@@ -97,6 +98,7 @@ CONSTFUNC int suggested_blank_lines(ast_t *first, ast_t *second) {
     case While:
     case For:
     case Block:
+    case Defer:
     case ConvertDef:
     case FunctionDef:
     case Lambda:
diff --git a/src/parse/controlflow.c b/src/parse/controlflow.c
index 74b3ea7a..1087e20e 100644
--- a/src/parse/controlflow.c
+++ b/src/parse/controlflow.c
@@ -79,6 +79,13 @@ ast_t *parse_pass(parse_ctx_t *ctx, const char *pos) {
     return match_word(&pos, "pass") ? NewAST(ctx->file, start, pos, Pass) : NULL;
 }
 
+ast_t *parse_defer(parse_ctx_t *ctx, const char *pos) {
+    const char *start = pos;
+    if (!match_word(&pos, "defer")) return NULL;
+    ast_t *body = expect(ctx, start, &pos, parse_block, "I expected a block to be deferred here");
+    return NewAST(ctx->file, start, pos, Defer, .body = body);
+}
+
 ast_t *parse_skip(parse_ctx_t *ctx, const char *pos) {
     const char *start = pos;
     if (!match_word(&pos, "continue") && !match_word(&pos, "skip")) return NULL;
diff --git a/src/parse/controlflow.h b/src/parse/controlflow.h
index 35a5a3f6..2ef093d4 100644
--- a/src/parse/controlflow.h
+++ b/src/parse/controlflow.h
@@ -5,6 +5,7 @@
 #include "context.h"
 
 ast_t *parse_block(parse_ctx_t *ctx, const char *pos);
+ast_t *parse_defer(parse_ctx_t *ctx, const char *pos);
 ast_t *parse_do(parse_ctx_t *ctx, const char *pos);
 ast_t *parse_for(parse_ctx_t *ctx, const char *pos);
 ast_t *parse_if(parse_ctx_t *ctx, const char *pos);
diff --git a/src/parse/expressions.c b/src/parse/expressions.c
index d6578b9e..3cb47669 100644
--- a/src/parse/expressions.c
+++ b/src/parse/expressions.c
@@ -191,9 +191,9 @@ ast_t *parse_term_no_suffix(parse_ctx_t *ctx, const char *pos) {
            || (term = parse_bool(ctx, pos)) || (term = parse_text(ctx, pos)) || (term = parse_path(ctx, pos))
            || (term = parse_lambda(ctx, pos)) || (term = parse_parens(ctx, pos)) || (term = parse_table(ctx, pos))
            || (term = parse_deserialize(ctx, pos)) || (term = parse_var(ctx, pos)) || (term = parse_list(ctx, pos))
-           || (term = parse_reduction(ctx, pos)) || (term = parse_pass(ctx, pos)) || (term = parse_skip(ctx, pos))
-           || (term = parse_stop(ctx, pos)) || (term = parse_return(ctx, pos)) || (term = parse_not(ctx, pos))
-           || (term = parse_inline_c(ctx, pos)));
+           || (term = parse_reduction(ctx, pos)) || (term = parse_pass(ctx, pos)) || (term = parse_defer(ctx, pos))
+           || (term = parse_skip(ctx, pos)) || (term = parse_stop(ctx, pos)) || (term = parse_return(ctx, pos))
+           || (term = parse_not(ctx, pos)) || (term = parse_inline_c(ctx, pos)));
     return term;
 }
 
diff --git a/src/parse/utils.c b/src/parse/utils.c
index b45e388e..2048a3ff 100644
--- a/src/parse/utils.c
+++ b/src/parse/utils.c
@@ -12,9 +12,10 @@
 #include "utils.h"
 
 static const char *keywords[] = {
-    "C_code", "_max_", "_min_", "and",  "assert", "break", "continue", "deserialize", "do",   "else",  "enum", "extend",
-    "for",    "func",  "if",    "in",   "lang",   "mod",   "mod1",     "no",          "none", "not",   "or",   "pass",
-    "return", "skip",  "skip",  "stop", "struct", "then",  "unless",   "use",         "when", "while", "xor",  "yes",
+    "C_code", "_max_",  "_min_",  "and",  "assert", "break",  "continue", "defer", "deserialize", "do",
+    "else",   "enum",   "extend", "for",  "func",   "if",     "in",       "lang",  "mod",         "mod1",
+    "no",     "none",   "not",    "or",   "pass",   "return", "skip",     "skip",  "stop",        "struct",
+    "then",   "unless", "use",    "when", "while",  "xor",    "yes",
 };
 
 CONSTFUNC bool is_keyword(const char *word) {
diff --git a/src/typecheck.c b/src/typecheck.c
index 0f720faa..3c820fa9 100644
--- a/src/typecheck.c
+++ b/src/typecheck.c
@@ -1123,7 +1123,8 @@ type_t *get_type(env_t *env, ast_t *ast) {
     case Skip: {
         return Type(AbortType);
     }
-    case Pass: return Type(VoidType);
+    case Pass:
+    case Defer: return Type(VoidType);
     case Negative: {
         ast_t *value = Match(ast, Negative)->value;
         type_t *t = get_type(env, value);
diff --git a/test/defer.tm b/test/defer.tm
new file mode 100644
index 00000000..8f0cb3be
--- /dev/null
+++ b/test/defer.tm
@@ -0,0 +1,79 @@
+func main()
+    x := 123
+    nums : @[Int]
+    do
+        defer
+            nums.insert(x)
+        x = 999
+
+    assert nums[] == [123]
+    assert x == 999
+
+    defer
+        say("All done!")
+
+    for word in ["first", "second", "third"]
+        defer
+            say("Got $word deferred")
+
+        if word == "second"
+            say("<skipped>")
+            skip
+        else if word == "third"
+            say("<stopped>")
+            stop
+
+        for i in 3
+            defer
+                say("Inner loop deferred $i")
+
+            if i == 2
+                say("<skipped inner>")
+                skip
+            else if i == 3
+                say("<stopped inner>")
+                stop
+
+            say("Made it through inner loop")
+
+        say("Made it through the loop")
+    
+    >> thunk := func(return_early=no)
+        say("Entering thunk")
+        defer
+            say("Deferred thunk cleanup")
+
+        if return_early
+            say("Returning early...")
+            return
+
+        say("Finished thunk")
+
+    >> thunk(no)
+    >> thunk(yes)
+
+    >> defer_func(yes)
+    >> defer_func(no)
+
+    >> counter := make_counter()
+    assert counter() == 1
+    assert counter() == 2
+    assert counter() == 3
+
+func defer_func(return_early=no)
+    say("Entering defer_func")
+    defer
+        say("Deferred defer_func cleanup")
+
+    if return_early
+        say("Returning early...")
+        return
+
+    say("Finished defer_func")
+
+func make_counter(->func(->Int))
+    i := 1
+    return func()
+        defer i += 1
+        return i
+
diff --git a/test/lambdas.tm b/test/lambdas.tm
index 5508dba2..1d1b2775 100644
--- a/test/lambdas.tm
+++ b/test/lambdas.tm
@@ -31,6 +31,6 @@ func main()
     fn := func()
         return func()
             return func()
-                say("$outer")
+                defer say("$outer")
                 return outer
     assert fn()()() == "Hello"