Locate and pass through captured locals to stateful closures

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

1 files changed, 165 insertions, 6 deletions

diff --git a/boot/lower.c b/boot/lower.c
index ba83b39..81f39a5 100644
--- a/boot/lower.c
+++ b/boot/lower.c
@@ -33,6 +33,18 @@ struct Local_Variable
     struct Type_Ref type; // TODO: type-check & fill out empty
 };
 
+// result of looking up a local that may turn out to be a capture for
+// the function we are currently inside.
+struct Lower_Local_Lookup_Capture
+{
+    // was the local found in an enclosing function?
+    bool captured;
+    // the captured local's declared type
+    struct Type_Ref type;
+    // which function does this capture belong to?
+    struct Function* captured_into;
+};
+
 // what kind of scope is this?
 // declared on every lexical scope so we can make sure that the scopes are balanced.
 enum Scope_Type
@@ -65,6 +77,14 @@ struct Scope
 {
     enum Scope_Type type;
     Array(struct Local_Variable) variables;
+
+    // function this scope is lexically contained in.
+    // for function scopes, references the function itself,
+    // for block scopes, references the enclosing function.
+    // used to check whether references cross function scope boundary,
+    // signifying a closure capture.
+    struct Function* containing_function;
+
     // TODO: add local closure functions & local types.
 };
 
@@ -668,6 +688,46 @@ lower_function_closure_type(struct Lower_Context* ctx, struct Function* fn)
     return ref.type_id;
 }
 
+// synthesize the per-lambda state struct that carries pointers to each
+// captured local.
+// each captured local is stored by reference.
+// NOTE(mel): the transpiler pass will output this as a normal structure,
+// does it make sense to distinguish these state types as something else?
+// it muddles distinguishing a type that the user wrote, and a type we synthesized.
+Type_Id
+lower_function_capture_state_type(struct Lower_Context* ctx, struct Function* fn)
+{
+    if (fn->capture_state_type_id != 0) return fn->capture_state_type_id;
+    if (array_length(&fn->captures) == 0) return 0;
+
+    Type_Id id = array_length(&ctx->unit->types.entries);
+    struct String name = lower_synthesize_type_name(ctx);
+    struct Type* type = type_new(id, TYPE_STRUCTURE, name, fn->ast_header->span);
+    type->synthetic = true;
+    type->depends_on = array_new(Type_Id, 8);
+    type->value.structure.fields = array_new(struct Field, 8);
+
+    FOR_EACH_ARRAY (struct Capture, cap, &fn->captures) {
+        // each field carries a pointer to the captured local. peel any
+        // existing modifiers off the captured type and prepend a reference.
+        // TODO: correctly nest all modifiers here!
+        struct Type_Ref ptr = {
+            .type_id = cap->type.type_id,
+            .mods = array_new(enum Type_Modifier, 4),
+        };
+        enum Type_Modifier ref_mod = TYPE_MOD_REFERENCE;
+        array_push(&ptr.mods, &ref_mod);
+        FOR_EACH_ARRAY (enum Type_Modifier, mod, &cap->type.mods) array_push(&ptr.mods, mod);
+
+        struct Field f = { .name = cap->name, .type = ptr };
+        array_push(&type->value.structure.fields, &f);
+    }
+
+    array_push(&ctx->unit->types.entries, &type);
+    fn->capture_state_type_id = id;
+    return id;
+}
+
 void
 lower_fill_type_body(struct Lower_Context* ctx, struct Type* type, struct Tree_Type* tree_type)
 {
@@ -858,6 +918,12 @@ struct Block* lower_block(struct Lower_Context* ctx, struct Tree_Block* tree_blo
 struct Statement* lower_statement(struct Lower_Context* ctx, struct Tree_Statement* tree_stmt);
 struct Expression* lower_expression(struct Lower_Context* ctx, struct Tree_Expression* tree_expr);
 bool lower_name_is_function_typed_local(struct Lower_Context* ctx, struct String name);
+struct Lower_Local_Lookup_Capture
+lower_local_lookup_capturing(struct Lower_Context* ctx, struct String name);
+void lower_record_capture(struct Function* fn, struct String name, struct Type_Ref type);
+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 Expression*
 lower_expression_integer_literal(struct Lower_Context* ctx, struct Tree_Expression* tree_expr)
@@ -890,16 +956,31 @@ lower_expression_boolean_literal(struct Lower_Context* ctx, struct Tree_Expressi
 struct Expression*
 lower_expression_name(struct Lower_Context* ctx, struct Tree_Expression* tree_expr)
 {
+    struct String name = tree_expr->value.name.name;
+
     // names that resolve to a catskill function become a function reference.
     Function_Id fn_id;
-    if (lower_function_lookup_by_name(ctx->unit, tree_expr->value.name.name, &fn_id)) {
+    if (lower_function_lookup_by_name(ctx->unit, name, &fn_id)) {
         // since this path is only taken outside direct call-positions,
         // the function will need a closure type to wrap it as a value.
         struct Function* fn = *array_at(struct Function*, &ctx->unit->functions.entries, fn_id);
         Type_Id closure_type_id = lower_function_closure_type(ctx, fn);
         return ir_make_function_ref(fn_id, closure_type_id, tree_expr->span);
     }
-    return ir_make_name(tree_expr->value.name.name, tree_expr->span);
+
+    // 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);
+        return ir_make_capture_ref(name, tree_expr->span);
+    }
+
+    // unresolved names will continue being bare names.
+    // this is required to account for all non-language objects that might
+    // come from c, or the runtime, or some c-header include we can't see.
+    // the backend will note any name reference which is still unresolved
+    // after taking every source into account.
+    return ir_make_name(name, tree_expr->span);
 }
 
 // any groups are discarded in the lowered representation, their presence
@@ -1239,9 +1320,22 @@ lower_expression_function(struct Lower_Context* ctx, struct Tree_Expression* tre
     fn->synthetic = true;
     lower_fill_function_signature(ctx, fn, tree_expr);
 
+    // lower the body inline so name lookups inside see the enclosing
+    // function's scope and can detect captures.
+    lower_push_scope_function(ctx, fn);
+    FOR_EACH_ARRAY (struct Param, param, &fn->params)
+        lower_declare_local(ctx, param->name, param->type);
+    fn->body = lower_block(ctx, fn->ast_body);
+    lower_pop_scope(ctx, SCOPE_TYPE_FUNCTION);
+
+    fn->completed = true; // mark the function as fully lowered.
+
     // a function literal is always in value context, we can synthesize
     // the closure type eagerly right here.
     Type_Id closure_type_id = lower_function_closure_type(ctx, fn);
+    // if the body captured anything, also synthesize the state type.
+    lower_function_capture_state_type(ctx, fn);
+
     return ir_make_function_ref(id, closure_type_id, tree_expr->span);
 }
 
@@ -1312,6 +1406,15 @@ lower_expression(struct Lower_Context* ctx, struct Tree_Expression* tree_expr)
     }
 }
 
+// look up the function whose body we are currently inside. (if there is one)
+struct Function*
+lower_current_function(struct Lower_Context* ctx)
+{
+    if (ctx->scope_stack.length == 0) return nil;
+    struct Scope* top = array_at(struct Scope, &ctx->scope_stack, ctx->scope_stack.length - 1);
+    return top->containing_function;
+}
+
 // push a fresh lexical scope onto the lowering stack.
 // do not forget to also pop this scope once we leave it!
 void
@@ -1320,6 +1423,21 @@ lower_push_scope(struct Lower_Context* ctx, enum Scope_Type type)
     struct Scope s = {
         .type = type,
         .variables = array_new(struct Local_Variable, 16),
+        // block scopes inherit the function from their upper scope
+        .containing_function = lower_current_function(ctx),
+    };
+    array_push(&ctx->scope_stack, &s);
+}
+
+// push a function scope.
+void
+lower_push_scope_function(struct Lower_Context* ctx, struct Function* owner)
+{
+    struct Scope s = {
+        .type = SCOPE_TYPE_FUNCTION,
+        .variables = array_new(struct Local_Variable, 16),
+        // the owner of this scope is the function that created it
+        .containing_function = owner,
     };
     array_push(&ctx->scope_stack, &s);
 }
@@ -1387,6 +1505,44 @@ lower_local_lookup(struct Lower_Context* ctx, struct String name, struct Type_Re
     return false;
 }
 
+// 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.
+struct Lower_Local_Lookup_Capture
+lower_local_lookup_capturing(struct Lower_Context* ctx, struct String name)
+{
+    struct Function* current = lower_current_function(ctx);
+
+    for (uint i = ctx->scope_stack.length; i > 0; --i) {
+        struct Scope* scope = array_at(struct Scope, &ctx->scope_stack, i - 1);
+        FOR_EACH_ARRAY (struct Local_Variable, var, &scope->variables) {
+            if (!string_equals(var->name, name)) continue;
+
+            bool from_enclosing = current != nil && scope->containing_function != current;
+            return (struct Lower_Local_Lookup_Capture){
+                .captured = from_enclosing,
+                .type = var->type,
+                .captured_into = from_enclosing ? current : nil,
+            };
+        }
+    }
+
+    // no capture.
+    return (struct Lower_Local_Lookup_Capture){ 0 };
+}
+
+// add a captured name to a function's capture list.
+void
+lower_record_capture(struct Function* fn, struct String name, struct Type_Ref type)
+{
+    // prevent duplication, closed-over locals are often referenced multiple times.
+    FOR_EACH_ARRAY (struct Capture, c, &fn->captures) {
+        if (string_equals(c->name, name)) return;
+    }
+    struct Capture cap = { .name = name, .type = type };
+    array_push(&fn->captures, &cap);
+}
+
 // is the named local a function-typed value?
 // (used to decide direct vs. indirect call dispatch when the call subject is a bare name.)
 bool
@@ -1735,18 +1891,21 @@ void
 lower_pass_2(struct Lower_Context* ctx)
 {
     FOR_EACH_ARRAY (struct Function*, fn, &ctx->unit->functions.entries) {
+        // skip functions which have already been processed.
+        // synthetic lambdas get lowered eagerly when encountered to gain
+        // access to the local stack at their location.
+        if ((*fn)->completed) continue;
         if (!(*fn)->ast_body) continue;
 
-        // push a function-level scope on top of the persistent top-level
+        // push a function-level scope on top of the persistent top-level unit
         // scope, holding the parameter names.
-        // the function body will also push a new block scope,
-        // separating the parameters and the local declarations.
-        lower_push_scope(ctx, SCOPE_TYPE_FUNCTION);
+        lower_push_scope_function(ctx, *fn);
         FOR_EACH_ARRAY (struct Param, param, &(*fn)->params) {
             lower_declare_local(ctx, param->name, param->type);
         }
 
         (*fn)->body = lower_block(ctx, (*fn)->ast_body);
+        (*fn)->completed = true;
 
         lower_pop_scope(ctx, SCOPE_TYPE_FUNCTION);
     }