Support multi-level closure captures

Signed-off-by: Mel <mel@rnrd.eu>

3 files changed, 45 insertions, 6 deletions

diff --git a/boot/ir.c b/boot/ir.c
index f86d63f..af631f8 100644
--- a/boot/ir.c
+++ b/boot/ir.c
@@ -468,6 +468,9 @@ struct Function
     // synthetic structure containing all captured locals. (only stateful lambdas)
     Type_Id capture_state_type_id;
 
+    // for nested functions this is the enclosing function. nil for top-level.
+    struct Function* parent_function;
+
     // tree-side header which produced this function node.
     // required for synthesizing the closure type after function is lowered.
     struct Tree_Function_Header* ast_header;
diff --git a/boot/lower.c b/boot/lower.c
index 81f39a5..4d00c5d 100644
--- a/boot/lower.c
+++ b/boot/lower.c
@@ -41,8 +41,8 @@ struct Lower_Local_Lookup_Capture
     bool captured;
     // the captured local's declared type
     struct Type_Ref type;
-    // which function does this capture belong to?
-    struct Function* captured_into;
+    // the function the local lives in.
+    struct Function* owner;
 };
 
 // what kind of scope is this?
@@ -924,6 +924,7 @@ void lower_record_capture(struct Function* fn, struct String name, struct Type_R
 void lower_push_scope_function(struct Lower_Context* ctx, struct Function* owner);
 void lower_pop_scope(struct Lower_Context* ctx, enum Scope_Type expected);
 void lower_declare_local(struct Lower_Context* ctx, struct String name, struct Type_Ref type);
+struct Function* lower_current_function(struct Lower_Context* ctx);
 
 struct Expression*
 lower_expression_integer_literal(struct Lower_Context* ctx, struct Tree_Expression* tree_expr)
@@ -971,7 +972,12 @@ lower_expression_name(struct Lower_Context* ctx, struct Tree_Expression* tree_ex
     // locals that were defined in an upper function scope are captured.
     struct Lower_Local_Lookup_Capture lookup = lower_local_lookup_capturing(ctx, name);
     if (lookup.captured) {
-        lower_record_capture(lookup.captured_into, name, lookup.type);
+        struct Function* fn = lower_current_function(ctx);
+        // we walk the enclosing function chain to record every in-between capture.
+        while (fn && fn != lookup.owner) {
+            lower_record_capture(fn, name, lookup.type);
+            fn = fn->parent_function;
+        }
         return ir_make_capture_ref(name, tree_expr->span);
     }
 
@@ -1318,6 +1324,7 @@ lower_expression_function(struct Lower_Context* ctx, struct Tree_Expression* tre
 
     struct Function* fn = *array_at(struct Function*, &ctx->unit->functions.entries, id);
     fn->synthetic = true;
+    fn->parent_function = lower_current_function(ctx);
     lower_fill_function_signature(ctx, fn, tree_expr);
 
     // lower the body inline so name lookups inside see the enclosing
@@ -1507,7 +1514,8 @@ lower_local_lookup(struct Lower_Context* ctx, struct String name, struct Type_Re
 
 // walk the scope stack looking for a local, which could be captured.
 // if the local was declared in some enclosing function that's not us,
-// the local becomes captured by our closure function.
+// the local becomes captured by our closure function (and any lambdas
+// nested between us and the local's owner).
 struct Lower_Local_Lookup_Capture
 lower_local_lookup_capturing(struct Lower_Context* ctx, struct String name)
 {
@@ -1522,7 +1530,7 @@ lower_local_lookup_capturing(struct Lower_Context* ctx, struct String name)
             return (struct Lower_Local_Lookup_Capture){
                 .captured = from_enclosing,
                 .type = var->type,
-                .captured_into = from_enclosing ? current : nil,
+                .owner = scope->containing_function,
             };
         }
     }
diff --git a/boot/transpile.c b/boot/transpile.c
index ac480cb..e0af3bf 100644
--- a/boot/transpile.c
+++ b/boot/transpile.c
@@ -92,6 +92,9 @@ struct Transpiler
     // expressions use this to decide whether to wrap themselves in parentheses
     // in the case where their parent expression is stronger than them.
     uint expression_parent_precedence;
+
+    // the function whose body we are in, if any.
+    struct Function* current_function;
 };
 
 void
@@ -389,6 +392,9 @@ transpile_visit_function_body(struct Visit* visit, struct Function* function)
         return;
     }
 
+    struct Function* previous_function = transpiler->current_function;
+    transpiler->current_function = function; // record current function context
+
     if (function->synthetic) {
         TRANSPILE_WRITE(" {\n");
         transpiler->indent_level++;
@@ -406,12 +412,14 @@ transpile_visit_function_body(struct Visit* visit, struct Function* function)
         transpiler->indent_level--;
         transpile_emit_indent(transpiler);
         TRANSPILE_WRITE("}\n");
+        transpiler->current_function = previous_function;
         return;
     }
 
     TRANSPILE_WRITE(" ");
     VISIT(visit_block, function->body);
     TRANSPILE_WRITE("\n");
+    transpiler->current_function = previous_function;
 
     if (function->closure_type_id != 0) transpile_function_emit_thunk(visit, function);
 }
@@ -511,7 +519,27 @@ transpile_visit_expression_reference(struct Visit* visit, struct Expression* exp
             bool first = true;
             FOR_EACH_ARRAY (struct Capture, cap, &fn->captures) {
                 if (!first) TRANSPILE_WRITE(", ");
-                TRANSPILE_WRITE(".%s = &%s", string_c_str(cap->name), string_c_str(cap->name));
+
+                bool from_constructors_state = false;
+                if (transpiler->current_function) {
+                    FOR_EACH_ARRAY (
+                        struct Capture, outer, &transpiler->current_function->captures) {
+                        if (string_equals(outer->name, cap->name)) {
+                            from_constructors_state = true;
+                            break;
+                        }
+                    }
+                }
+
+                if (from_constructors_state) {
+                    // we are multiple closures deep, steal the reference
+                    // from our closure parent
+                    TRANSPILE_WRITE(
+                        ".%s = state->%s", string_c_str(cap->name), string_c_str(cap->name));
+                } else {
+                    // capture reference to a local on the stack
+                    TRANSPILE_WRITE(".%s = &%s", string_c_str(cap->name), string_c_str(cap->name));
+                }
                 first = false;
             }
             TRANSPILE_WRITE(" }");