Implement closures for function declarations

src/vstc/capture_finder.rs

@@ -22,20 +22,31 @@ impl CaptureFinder {
       return;
     }
 
+    let mut insert = |n: &String| {
+      let inserted = self.names.insert(n.clone());
+
+      if inserted {
+        self.ordered_names.push(n.clone());
+      }
+    };
+
     match self.outside_scope.get(&name) {
       None => std::panic!("Unresolved name"),
       Some(MappedName::Definition(_)) => {}, // Not capture - just definition
-      Some(MappedName::Register(_)) => {
-        let inserted = self.names.insert(name.clone());
+      Some(MappedName::Register(_)) => insert(&name),
+      Some(MappedName::QueuedFunction(qfn)) => {
+        for cap in &qfn.capture_params {
+          if scope.get(cap).is_some() {
+            std::panic!("Not implemented: Nested capture edge case");
+          }
 
-        if inserted {
-          self.ordered_names.push(name);
+          insert(cap);
         }
       },
     }
   }
 
-  fn fn_decl(&mut self, parent_scope: &Scope, decl: &swc_ecma_ast::FnDecl) {
+  pub fn fn_decl(&mut self, parent_scope: &Scope, decl: &swc_ecma_ast::FnDecl) {
     let scope = parent_scope.nest();
 
     scope.set(

src/vstc/expression_compiler.rs

@@ -204,6 +204,7 @@ impl<'a> ExpressionCompiler<'a> {
               swc_ecma_ast::Expr::Ident(ident) => match ec.scope.get(&ident.sym.to_string()) {
                 None => std::panic!("Unresolved identifier"),
                 Some(MappedName::Definition(_)) => std::panic!("Invalid: definition mutation"),
+                Some(MappedName::QueuedFunction(_)) => std::panic!("Invalid: assign to declaration"),
                 Some(MappedName::Register(reg)) => AssignTarget::Register(reg),
               },
               swc_ecma_ast::Expr::This(_) => AssignTarget::Register("this".to_string()),
@@ -220,6 +221,7 @@ impl<'a> ExpressionCompiler<'a> {
             return match ec.scope.get(&ident.sym.to_string()) {
               None => std::panic!("Unresolved identifier"),
               Some(MappedName::Definition(_)) => std::panic!("Invalid: definition mutation"),
+              Some(MappedName::QueuedFunction(_)) => std::panic!("Invalid: assign to declaration"),
               Some(MappedName::Register(reg)) => AssignTarget::Register(reg),
             };
           }
@@ -299,6 +301,7 @@ impl<'a> ExpressionCompiler<'a> {
               match self.scope.get(&ident.id.sym.to_string()) {
                 None => std::panic!("Unresolved identifier"),
                 Some(MappedName::Definition(_)) => std::panic!("Invalid: definition mutation"),
+                Some(MappedName::QueuedFunction(_)) => std::panic!("Invalid: assign to declaration"),
                 Some(MappedName::Register(reg)) => reg,
               }
             ),
@@ -742,7 +745,7 @@ impl<'a> ExpressionCompiler<'a> {
     self.fnc.queue.add(QueuedFunction {
       definition_name: definition_name.clone(),
       fn_name: fn_name.clone(),
-      extra_params: cf.ordered_names.clone(),
+      capture_params: cf.ordered_names.clone(),
       function: fn_.function.clone(),
     }).expect("Failed to queue function");
 
@@ -753,7 +756,23 @@ impl<'a> ExpressionCompiler<'a> {
       );
     }
 
+    return self.capturing_fn_ref(
+      fn_name,
+      &definition_name,
+      &cf.ordered_names,
+      target_register,
+    );
+  }
+
+  pub fn capturing_fn_ref(
+    &mut self,
+    fn_name: Option<String>,
+    definition_name: &String,
+    captures: &Vec<String>,
+    target_register: Option<String>,
+  ) -> CompiledExpression {
     let mut nested_registers = Vec::<String>::new();
+    let mut sub_nested_registers = Vec::<String>::new();
 
     let reg = match target_register {
       None => {
@@ -771,23 +790,39 @@ impl<'a> ExpressionCompiler<'a> {
 
     let mut bind_instr = format!("  bind @{} [", definition_name);
 
-    for i in 0..cf.ordered_names.len() {
-      let captured_name = &cf.ordered_names[i];
+    for i in 0..captures.len() {
+      let captured_name = &captures[i];
 
       if i > 0 {
         bind_instr += ", ";
       }
 
       bind_instr += &match self.scope.get(captured_name) {
-        None => std::panic!(""),
-        Some(MappedName::Definition(_)) => std::panic!(""),
+        None => std::panic!("Captured names should always be in scope"),
+        Some(MappedName::Definition(_)) => std::panic!("Definitions should never be recorded as captures"),
         Some(MappedName::Register(cap_reg)) => format!("%{}", cap_reg),
+        Some(MappedName::QueuedFunction(qfn)) => {
+          let mut compiled_ref = self.capturing_fn_ref(
+            qfn.fn_name.clone(),
+            &qfn.definition_name,
+            &qfn.capture_params,
+            None,
+          );
+
+          sub_nested_registers.append(&mut compiled_ref.nested_registers);
+
+          compiled_ref.value_assembly
+        },
       };
     }
 
     bind_instr += &format!("] %{}", reg);
     self.fnc.definition.push(bind_instr);
 
+    for reg in sub_nested_registers {
+      self.fnc.reg_allocator.release(&reg);
+    }
+
     return CompiledExpression {
       value_assembly: format!("%{}", reg),
       nested_registers: nested_registers,
@@ -840,12 +875,16 @@ impl<'a> ExpressionCompiler<'a> {
 
     let mapped = self.scope.get(&ident_string).expect("Identifier not found in scope");
 
-    let value_assembly = match mapped {
-      MappedName::Register(reg) => "%".to_string() + &reg,
-      MappedName::Definition(def) => "@".to_string() + &def,
+    return match mapped {
+      MappedName::Register(reg) => self.inline("%".to_string() + &reg, target_register),
+      MappedName::Definition(def) => self.inline("@".to_string() + &def, target_register),
+      MappedName::QueuedFunction(qfn) => self.capturing_fn_ref(
+        qfn.fn_name.clone(),
+        &qfn.definition_name,
+        &qfn.capture_params,
+        target_register,
+      ),
     };
-
-    return self.inline(value_assembly, target_register);
   }
 }
 
@@ -951,6 +990,7 @@ impl TargetAccessor {
       Ident(ident) => match ec.scope.get(&ident.sym.to_string()) {
         None => std::panic!("Unresolved identifier"),
         Some(MappedName::Definition(_)) => std::panic!("Invalid: definition mutation"),
+        Some(MappedName::QueuedFunction(_)) => std::panic!("Invalid: assign to declaration"),
         Some(MappedName::Register(reg)) => TargetAccessor::Register(reg),
       },
       This(_) => TargetAccessor::Register("this".to_string()),

src/vstc/function_compiler.rs

@@ -1,16 +1,18 @@
 use std::rc::Rc;
 use std::cell::RefCell;
+use std::collections::{HashMap, HashSet};
 use queues::*;
 
 use super::name_allocator::NameAllocator;
 use super::expression_compiler::ExpressionCompiler;
-use super::scope::{Scope, MappedName, ScopeTrait};
+use super::scope::{Scope, MappedName, ScopeTrait, init_scope};
+use super::capture_finder::CaptureFinder;
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct QueuedFunction {
   pub definition_name: String,
   pub fn_name: Option<String>,
-  pub extra_params: Vec<String>,
+  pub capture_params: Vec<String>,
   pub function: swc_ecma_ast::Function,
 }
 
@@ -56,7 +58,7 @@ impl FunctionCompiler {
     self_.queue.add(QueuedFunction {
       definition_name: definition_name.clone(),
       fn_name: fn_name,
-      extra_params: Vec::new(),
+      capture_params: Vec::new(),
       function: fn_.clone(),
     }).expect("Failed to queue function");
 
@@ -65,7 +67,7 @@ impl FunctionCompiler {
         Ok(qfn) => self_.compile_fn(
           qfn.definition_name,
           qfn.fn_name,
-          qfn.extra_params,
+          qfn.capture_params,
           &qfn.function,
           parent_scope,
         ),
@@ -80,7 +82,7 @@ impl FunctionCompiler {
     &mut self,
     definition_name: String,
     fn_name: Option<String>,
-    mut extra_params: Vec<String>,
+    mut capture_params: Vec<String>,
     fn_: &swc_ecma_ast::Function,
     parent_scope: &Scope,
   ) {
@@ -103,7 +105,7 @@ impl FunctionCompiler {
     heading += " = function(";
 
     let mut params = Vec::<String>::new();
-    params.append(&mut extra_params);
+    params.append(&mut capture_params);
 
     for p in &fn_.params {
       match &p.pat {
@@ -254,6 +256,8 @@ impl FunctionCompiler {
     block: &swc_ecma_ast::BlockStmt,
     scope: &Scope,
   ) {
+    let mut function_decls = Vec::<swc_ecma_ast::FnDecl>::new();
+
     for statement in &block.stmts {
       use swc_ecma_ast::Stmt::*;
 
@@ -280,27 +284,7 @@ impl FunctionCompiler {
   
           match decl {
             Class(_) => std::panic!("Not implemented: Class declaration"),
-            Fn(fn_) => {
-              let fn_name = fn_.ident.sym.to_string();
-
-              let definition_name = self
-                .definition_allocator
-                .borrow_mut()
-                .allocate(&fn_name)
-              ;
-
-              scope.set(
-                fn_name.clone(),
-                MappedName::Definition(definition_name.clone()),
-              );
-
-              self.queue.add(QueuedFunction {
-                definition_name: definition_name,
-                fn_name: Some(fn_name),
-                extra_params: Vec::new(),
-                function: fn_.function.clone(),
-              }).expect("Failed to queue function");
-            },
+            Fn(fn_) => function_decls.push(fn_.clone()),
             Var(var_decl) => self.populate_block_scope_var_decl(var_decl, scope),
             TsInterface(_) => {},
             TsTypeAlias(_) => {},
@@ -311,6 +295,89 @@ impl FunctionCompiler {
         Expr(_) => {},
       };
     }
+
+    // Create a synth scope where the function decls that can co-mingle are
+    // present but don't signal any nested captures. This allows us to first
+    // construct all the direct captures and use that to find the complete
+    // captures.
+    let synth_scope = scope.nest();
+
+    for fn_ in &function_decls {
+      synth_scope.set(
+        fn_.ident.sym.to_string(),
+        MappedName::Register("".to_string()),
+      );
+    }
+
+    let mut direct_captures_map = HashMap::<String, Vec<String>>::new();
+
+    for fn_ in &function_decls {
+      let mut cf = CaptureFinder::new(synth_scope.clone());
+      cf.fn_decl(&init_scope(), fn_);
+
+      direct_captures_map.insert(
+        fn_.ident.sym.to_string(),
+        cf.ordered_names,
+      );
+    }
+
+    for fn_ in &function_decls {
+      let mut full_captures = Vec::<String>::new();
+      let mut full_captures_set = HashSet::<String>::new();
+
+      let mut cap_queue = Queue::<String>::new();
+
+      for dc in direct_captures_map.get(&fn_.ident.sym.to_string())
+        .expect("direct captures not found")
+      {
+        cap_queue.add(dc.clone())
+          .expect("Failed to add to queue");
+      }
+
+      loop {
+        let cap = match cap_queue.remove() {
+          Ok(c) => c,
+          Err(_) => { break; },
+        };
+
+        let is_new = full_captures_set.insert(cap.clone());
+
+        if !is_new {
+          continue;
+        }
+
+        full_captures.push(cap.clone());
+
+        for nested_caps in direct_captures_map.get(&cap) {
+          for nested_cap in nested_caps {
+            cap_queue.add(nested_cap.clone())
+              .expect("Failed to add to queue");
+          }
+        }
+      }
+
+      let fn_name = fn_.ident.sym.to_string();
+
+      let definition_name = self
+        .definition_allocator
+        .borrow_mut()
+        .allocate(&fn_name)
+      ;
+
+      let qf = QueuedFunction {
+        definition_name: definition_name,
+        fn_name: Some(fn_name.clone()),
+        capture_params: full_captures,
+        function: fn_.function.clone(),
+      };
+
+      scope.set(
+        fn_name.clone(),
+        MappedName::QueuedFunction(qf.clone()),
+      );
+
+      self.queue.add(qf).expect("Failed to queue function");
+    }
   }
 
   fn populate_block_scope_var_decl(
@@ -360,6 +427,7 @@ impl FunctionCompiler {
               self.reg_allocator.release(reg);
             },
             MappedName::Definition(_) => {},
+            MappedName::QueuedFunction(_) => {},
           }
         }
       },

src/vstc/scope.rs

@@ -2,10 +2,13 @@ use std::collections::HashMap;
 use std::rc::Rc;
 use std::cell::RefCell;
 
+use super::function_compiler::QueuedFunction;
+
 #[derive(Clone, Debug)]
 pub enum MappedName {
   Register(String),
   Definition(String),
+  QueuedFunction(QueuedFunction),
 }
 
 pub struct ScopeData {