Support multi-level closure captures

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

1 files changed, 29 insertions, 1 deletions

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(" }");