Locate and pass through captured locals to stateful closures

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

1 files changed, 75 insertions, 41 deletions

diff --git a/boot/transpile.c b/boot/transpile.c
index 9b8f420..ac480cb 100644
--- a/boot/transpile.c
+++ b/boot/transpile.c
@@ -331,12 +331,15 @@ transpile_function_emit_signature(struct Visit* visit, struct Function* function
 {
     TRANSPILE_PREAMBLE
 
+    bool has_captures = function->capture_state_type_id != 0;
+
     VISIT(visit_type_ref, &function->return_type);
     TRANSPILE_WRITE(" %s(", transpile_function_c_name(function));
     bool need_separator = false;
     if (function->synthetic) {
-        // synthetic functions implicitly take in state
-        TRANSPILE_WRITE("void* state");
+        // synthetic functions implicitly take in state.
+        // TODO(mel): can we avoid state_void?
+        TRANSPILE_WRITE("void* %s", has_captures ? "state_void" : "state");
         need_separator = true;
     }
 
@@ -386,13 +389,18 @@ transpile_visit_function_body(struct Visit* visit, struct Function* function)
         return;
     }
 
-    // add a void cast for the implicit state argument, in case synthetic
-    // function does not need to use it
     if (function->synthetic) {
         TRANSPILE_WRITE(" {\n");
         transpiler->indent_level++;
         transpile_emit_indent(transpiler);
-        TRANSPILE_WRITE("(void)state;\n");
+        if (function->capture_state_type_id != 0) {
+            // TODO(mel): can we avoid this?
+            struct Type* state_type =
+                transpile_type_at(transpiler, function->capture_state_type_id);
+            TRANSPILE_WRITE("struct %s* state = state_void;\n", string_c_str(state_type->name));
+        } else {
+            TRANSPILE_WRITE("(void)state;\n");
+        }
         FOR_EACH_ARRAY (struct Statement*, statement, &function->body->statements)
             VISIT(visit_statement, *statement);
         transpiler->indent_level--;
@@ -459,11 +467,65 @@ transpile_visit_expression_boolean_literal(struct Visit* visit, struct Expressio
     TRANSPILE_WRITE("%s", expression->value.bool_literal.value ? "true" : "false");
 }
 
+// emit a reference.
+// the reference's inner kind. callers in subject-of-call position
+// short-circuit FUNCTION references before reaching this handler.
 void
-transpile_visit_expression_name(struct Visit* visit, struct Expression* expression)
+transpile_visit_expression_reference(struct Visit* visit, struct Expression* expression)
 {
     TRANSPILE_PREAMBLE
-    TRANSPILE_WRITE("%s", string_c_str(expression->value.name.name));
+    struct Expression_Reference* ref = &expression->value.reference;
+    switch (ref->kind) {
+    case REFERENCE_NAME:
+        TRANSPILE_WRITE("%s", string_c_str(ref->value.name.name));
+        return;
+
+    case REFERENCE_CAPTURE:
+        TRANSPILE_WRITE("(*state->%s)", string_c_str(ref->value.capture.name));
+        return;
+
+    case REFERENCE_FUNCTION: {
+        struct Function* fn = *array_at(
+            struct Function*, &transpiler->unit->functions.entries,
+            ref->value.function.function_id);
+        struct Type* closure_type =
+            transpile_type_at(transpiler, ref->value.function.closure_type_id);
+
+        const ascii* call_name;
+        struct String thunk_name = string_empty();
+        if (fn->synthetic) {
+            // synthetic lambdas already carry the fat signature
+            call_name = transpile_function_c_name(fn);
+        } else {
+            // named functions go through their thunk converting their
+            // thin calling convention, to the fat lambda convention.
+            thunk_name = transpile_function_thunk_name(fn);
+            call_name = string_c_str(thunk_name);
+        }
+
+        TRANSPILE_WRITE("(struct %s){ .state = ", string_c_str(closure_type->name));
+        if (fn->capture_state_type_id != 0) {
+            // closure with closed-over state, construct a state value.
+            struct Type* state_type = transpile_type_at(transpiler, fn->capture_state_type_id);
+            TRANSPILE_WRITE("&(struct %s){ ", string_c_str(state_type->name));
+            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));
+                first = false;
+            }
+            TRANSPILE_WRITE(" }");
+        } else {
+            TRANSPILE_WRITE("nil");
+        }
+        TRANSPILE_WRITE(", .call = %s }", call_name);
+        return;
+    }
+
+    case REFERENCE_NONE:
+        TRANSPILE_WRITE("/* invalid reference */");
+        return;
+    }
 }
 
 void
@@ -533,11 +595,14 @@ transpile_visit_expression_call(struct Visit* visit, struct Expression* expressi
     struct Expression* subject = expression->value.call.subject;
     bool is_indirect = expression->value.call.is_indirect;
 
-    if (!is_indirect && subject && subject->kind == EXPRESSION_FUNCTION_REF) {
+    bool subject_is_function_ref =
+        subject && subject->kind == EXPRESSION_REFERENCE
+        && subject->value.reference.kind == REFERENCE_FUNCTION;
+    if (!is_indirect && subject_is_function_ref) {
         // direct calls emit the subject bare
         struct Function* fn = *array_at(
             struct Function*, &transpiler->unit->functions.entries,
-            subject->value.function_ref.function_id);
+            subject->value.reference.value.function.function_id);
         TRANSPILE_WRITE("%s", transpile_function_c_name(fn));
     } else if (is_indirect) {
         // route through the fat closure value's call slot, passing state
@@ -620,36 +685,6 @@ transpile_visit_expression_construct(struct Visit* visit, struct Expression* exp
 }
 
 void
-transpile_visit_expression_function_ref(struct Visit* visit, struct Expression* expression)
-{
-    TRANSPILE_PREAMBLE
-
-    struct Function* fn = *array_at(
-        struct Function*, &transpiler->unit->functions.entries,
-        expression->value.function_ref.function_id);
-
-    // we only get here when the function reference is explicitly a *value*,
-    // any direct calls are filtered out earlier.
-    struct Type* closure_type =
-        transpile_type_at(transpiler, expression->value.function_ref.closure_type_id);
-
-    const ascii* call_name;
-    struct String thunk_name = string_empty();
-    if (fn->synthetic) {
-        // synthetic functions are already fat-by-default, can call directly
-        call_name = transpile_function_c_name(fn);
-    } else {
-        // thin named functions go through their generated thunk to match the
-        // fat calling convention
-        thunk_name = transpile_function_thunk_name(fn);
-        call_name = string_c_str(thunk_name);
-    }
-
-    TRANSPILE_WRITE(
-        "(struct %s){ .state = nil, .call = %s }", string_c_str(closure_type->name), call_name);
-}
-
-void
 transpile_visit_expression_increment_decrement(struct Visit* visit, struct Expression* expression)
 {
     TRANSPILE_PREAMBLE
@@ -839,7 +874,7 @@ struct Visit_Table transpile_body_visit_functions = {
     .visit_expression_float_literal = transpile_visit_expression_float_literal,
     .visit_expression_string_literal = transpile_visit_expression_string_literal,
     .visit_expression_boolean_literal = transpile_visit_expression_boolean_literal,
-    .visit_expression_name = transpile_visit_expression_name,
+    .visit_expression_reference = transpile_visit_expression_reference,
     .visit_expression_unary_operation = transpile_visit_expression_unary_operation,
     .visit_expression_binary_operation = transpile_visit_expression_binary_operation,
     .visit_expression_sizeof_operation = transpile_visit_expression_sizeof_operation,
@@ -849,7 +884,6 @@ struct Visit_Table transpile_body_visit_functions = {
     .visit_expression_cast = transpile_visit_expression_cast,
     .visit_expression_construct = transpile_visit_expression_construct,
     .visit_expression_increment_decrement = transpile_visit_expression_increment_decrement,
-    .visit_expression_function_ref = transpile_visit_expression_function_ref,
 };
 
 // walk a lowered translation unit and emit c source into the transpiler's output.