Propagate closures across multiple function calls

src/compiler.cpp

@@ -72,6 +72,13 @@ struct Context {
     return i;
   }
 
+  template <class T>
+  Result<int64_t> add_const(const T& val)
+    requires std::derived_from<T, Object>
+  {
+    return add_const(TRY(val.copy_value()));
+  }
+
   Result<int64_t> add_var(const Value& sym) {
     auto idx = variables_dict.get(sym);
     if (idx.has_value()) {
@@ -118,7 +125,12 @@ struct Context {
 
     if (!parent) return ERROR(KeyError);
 
-    TRY(parent->get_var(sym));
+    auto maybe_var = parent->get_var(sym);
+    // If parent has no variable with name <sym>, then we need to propagate the
+    // closure to it. The variable may be defined somewhere in outer scope.
+    if (!maybe_var.has_value()) {
+      TRY(parent->get_closure(sym));
+    }
 
     int64_t i = closures.size();
     closures = TRY(closures.append(sym));
@@ -639,9 +651,29 @@ Result<Expression> Compiler::compile_fn(Context& context, Symbol& op,
   for (uint64_t i = 0; i < ctx.closures.size(); i++) {
     Value sym = TRY(ctx.closures.get(i));
 
-    int64_t var_reg = TRY(context.get_var(sym));
-    uint64_t vr = context.alloc_reg();
-    TRY(ex_res.add_opcode(Oc::Mov, {0, (int64_t)vr}, {0, (int64_t)var_reg}));
+    // If this is a regular variable from outer scope - it is present in a
+    // register, and as such we can just move it to the correct place
+    auto maybe_var_reg = context.get_var(sym);
+    if (maybe_var_reg.has_value()) {
+      int64_t var_reg = maybe_var_reg.value();
+      uint64_t vr = context.alloc_reg();
+      TRY(ex_res.add_opcode(Oc::Mov, {0, (int64_t)vr}, {0, (int64_t)var_reg}));
+      continue;
+    }
+
+    // Otherwise this can be a variable we get from the intermediate closure
+    // (such as in case of 3 nested functions where a var we close over is
+    // defined in the outermost function.
+    auto maybe_closure = context.get_closure(sym);
+    if (maybe_closure.has_value()) {
+      int64_t closure_id = maybe_closure.value();
+      uint64_t vr = context.alloc_reg();
+      TRY(ex_res.add_opcode(Oc::ClosureLoad, {0, (int64_t)vr},
+                            {0, (int64_t)closure_id}));
+      continue;
+    }
+
+    return ERROR(NotImplemented);
   }
 
   TRY(ex_res.add_opcode(Oc::MakeClosure, {0, (int64_t)reg},
@@ -849,7 +881,7 @@ Result<Expression> Compiler::compile_symbol(Context& context,
 
   if (!maybe_stdlib_fun.has_error()) {
     auto stdlib_fun = TRY(StdlibFunction::create(maybe_stdlib_fun.value()));
-    int64_t c = TRY(context.add_const(TRY(stdlib_fun.copy_value())));
+    int64_t c = TRY(context.add_const(stdlib_fun));
 
     uint64_t reg = context.alloc_reg();
     TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)reg}, {1, (int64_t)c}));

test/function.vli

@@ -12,7 +12,17 @@
   (assert (= (square 4) 16))
   )
 
+;; Closures should work at least across one scope
 (let ((x 42))
   (assert (= ((fn (y) (+ x y) 1)
               43)))
   )
+
+;; Closures should work across multiple scopes
+(let ((x 42))
+  (assert (=
+           (((fn (y)
+                 (fn () (+ x y)))
+             1))
+           43))
+  )