feat: Add lambda collection and specialization

Introduces `LambdaCollector` to gather lambda functions and populate the
function registry. This enables the `Specializer` to work with
user-defined functions.

The `Environment` struct is updated to manage the `function_registry`
and `monomorph_cache`, which are essential for the specialization
process.

The `link` method in `Environment` now incorporates lambda collection
and node specialization before applying TCO optimization. This ensures
that lambdas are properly processed and specialized for potential
performance gains.

The `Specializer`'s `new` constructor has been modified to accept and
initialize the `MonoCache` through an `Rc<RefCell<MonoCache>>`. This
allows the cache to be shared across different specialized functions.

Also includes minor refactoring and type adjustments in `specializer.rs`
for better clarity and consistency.

src/ast/compiler/mod.rs

@@ -6,6 +6,7 @@ pub mod dumper;
 pub mod macros;
 pub mod type_checker;
 pub mod specializer;
+pub mod lambda_collector;
 
 pub use binder::*;
 pub use bound_nodes::*;

src/ast/compiler/specializer.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 use std::rc::Rc;
 use std::cell::RefCell;
 use crate::ast::types::{StaticType, Value, Signature};
-use crate::ast::compiler::bound_nodes::{BoundKind, Address, TypedNode, BoundNode};
+use crate::ast::compiler::bound_nodes::{BoundKind, Address, TypedNode};
 use crate::ast::nodes::Node;
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -11,15 +11,17 @@ pub struct MonoCacheKey {
     pub arg_types: Vec<StaticType>,
 }
 
-pub type CompileFunc = Rc<dyn Fn(Rc<BoundNode>, &[StaticType]) -> Result<(Value, StaticType), String>>;
+pub type CompileFunc = Rc<dyn Fn(TypedNode, &[StaticType]) -> Result<(Value, StaticType), String>>;
 pub type RtlLookupFunc = Rc<dyn Fn(&str, &[StaticType]) -> Option<(Value, StaticType)>>;
 
 pub trait FunctionRegistry {
-    fn resolve(&self, addr: Address) -> Option<Rc<BoundNode>>;
+    fn resolve(&self, addr: Address) -> Option<TypedNode>;
 }
 
+pub type MonoCache = HashMap<MonoCacheKey, (Value, StaticType)>;
+
 pub struct Specializer {
-    cache: RefCell<HashMap<MonoCacheKey, (Value, StaticType)>>,
+    pub cache: Rc<RefCell<MonoCache>>,
     registry: Option<Rc<dyn FunctionRegistry>>,
     compiler: Option<CompileFunc>,
     rtl_lookup: Option<RtlLookupFunc>,
@@ -30,9 +32,10 @@ impl Specializer {
         registry: Option<Rc<dyn FunctionRegistry>>, 
         compiler: Option<CompileFunc>,
         rtl_lookup: Option<RtlLookupFunc>,
+        cache: Option<Rc<RefCell<MonoCache>>>,
     ) -> Self {
         Self {
-            cache: RefCell::new(HashMap::new()),
+            cache: cache.unwrap_or_else(|| Rc::new(RefCell::new(HashMap::new()))),
             registry,
             compiler,
             rtl_lookup,
@@ -208,7 +211,7 @@ impl Specializer {
 mod tests {
     use super::*;
     use crate::ast::types::{Identity, NodeIdentity, SourceLocation, StaticType, Value, Signature};
-    use crate::ast::compiler::bound_nodes::{BoundKind, Address, BoundNode, TypedNode};
+    use crate::ast::compiler::bound_nodes::{BoundKind, Address, TypedNode};
     use crate::ast::nodes::Symbol;
     use std::rc::Rc;
 
@@ -226,7 +229,7 @@ mod tests {
 
     // Mock Registry
     struct MockRegistry {
-        functions: HashMap<Address, Rc<BoundNode>>,
+        functions: HashMap<Address, TypedNode>,
     }
     
     impl MockRegistry {
@@ -234,13 +237,13 @@ mod tests {
             Self { functions: HashMap::new() }
         }
         
-        fn register(&mut self, addr: Address, node: BoundNode) {
-            self.functions.insert(addr, Rc::new(node));
+        fn register(&mut self, addr: Address, node: TypedNode) {
+            self.functions.insert(addr, node);
         }
     }
     
     impl FunctionRegistry for MockRegistry {
-        fn resolve(&self, addr: Address) -> Option<Rc<BoundNode>> {
+        fn resolve(&self, addr: Address) -> Option<TypedNode> {
             self.functions.get(&addr).cloned()
         }
     }
@@ -253,24 +256,24 @@ mod tests {
         let name = Symbol::from("test_func");
         
         // Def: (fn [x] x) -- generic identity
-        let func_node = BoundNode {
+        let func_node = TypedNode {
             identity: make_identity(),
             kind: BoundKind::Lambda { 
                 param_count: 1, 
                 upvalues: vec![], 
-                body: Rc::new(BoundNode { identity: make_identity(), kind: BoundKind::Nop, ty: () }) 
+                body: Rc::new(TypedNode { identity: make_identity(), kind: BoundKind::Nop, ty: StaticType::Void }) 
             },
-            ty: ()
+            ty: StaticType::Void
         };
         registry.register(addr, func_node);
 
         // Setup Compiler Mock
-        let compiler: CompileFunc = Rc::new(|_node: Rc<BoundNode>, _args: &[StaticType]| -> Result<(Value, StaticType), String> {
+        let compiler: CompileFunc = Rc::new(|_node: TypedNode, _args: &[StaticType]| -> Result<(Value, StaticType), String> {
             // Return a specialized "compiled" value
             Ok((Value::Int(12345), StaticType::Int))
         });
 
-        let spec = Specializer::new(Some(Rc::new(registry)), Some(compiler), None);
+        let spec = Specializer::new(Some(Rc::new(registry)), Some(compiler), None, None);
         
         // Call(Get(Local(0)), [Arg(Int)])
         let callee = make_typed_node(BoundKind::Get { addr, name: name.clone() }, StaticType::Any);
@@ -300,7 +303,7 @@ mod tests {
 
     #[test]
     fn test_specialize_skips_unknown_types() {
-        let spec = Specializer::new(None, None, None);
+        let spec = Specializer::new(None, None, None, None);
         let name0 = Symbol::from("f");
         let name1 = Symbol::from("x");
         
@@ -330,7 +333,7 @@ mod tests {
     #[test]
     fn test_specialize_uses_cache() {
         // Setup cache with a pre-specialized function for (Int) -> Int
-        let spec = Specializer::new(None, None, None);
+        let spec = Specializer::new(None, None, None, None);
         
         let addr = Address::Local(0);
         let name = Symbol::from("cached_func");
@@ -380,7 +383,7 @@ mod tests {
             }
         });
         
-        let spec = Specializer::new(None, None, Some(rtl_lookup));
+        let spec = Specializer::new(None, None, Some(rtl_lookup), None);
         
         let addr = Address::Global(10);
         let name = Symbol::from("rtl_func");
@@ -411,7 +414,7 @@ mod tests {
 
     #[test]
     fn test_specialize_preserves_tail_call() {
-        let spec = Specializer::new(None, None, None);
+        let spec = Specializer::new(None, None, None, None);
         
         let addr = Address::Local(0);
         let name = Symbol::from("tail_func");

src/ast/environment.rs

@@ -9,22 +9,44 @@ use crate::ast::compiler::{TypedNode, TypeChecker};
 use crate::ast::vm::{VM, TracingObserver};
 
 use crate::ast::compiler::tco::TCO;
+use crate::ast::compiler::lambda_collector::LambdaCollector;
 use crate::ast::compiler::dumper::Dumper;
 use crate::ast::compiler::macros::{MacroExpander, MacroRegistry, MacroEvaluator};
+use crate::ast::compiler::specializer::{Specializer, MonoCache, FunctionRegistry};
 use crate::ast::rtl;
+use crate::ast::rtl::intrinsics;
+use crate::ast::compiler::bound_nodes::{BoundKind, Address};
 
 pub struct Environment {
     pub global_names: Rc<RefCell<HashMap<Symbol, u32>>>,
     pub global_types: Rc<RefCell<HashMap<u32, StaticType>>>,
     pub global_values: Rc<RefCell<Vec<Value>>>,
+    pub function_registry: Rc<RefCell<HashMap<u32, TypedNode>>>,
+    pub monomorph_cache: Rc<RefCell<MonoCache>>,
     pub debug_mode: bool,
 }
 
+struct EnvFunctionRegistry {
+    registry: Rc<RefCell<HashMap<u32, TypedNode>>>,
+}
+
+impl FunctionRegistry for EnvFunctionRegistry {
+    fn resolve(&self, addr: Address) -> Option<TypedNode> {
+        if let Address::Global(idx) = addr {
+            self.registry.borrow().get(&idx).cloned()
+        } else {
+            None
+        }
+    }
+}
+
+
 /// Evaluator used during macro expansion to allow compile-time logic.
 struct RuntimeMacroEvaluator {
     global_names: Rc<RefCell<HashMap<Symbol, u32>>>,
     global_types: Rc<RefCell<HashMap<u32, StaticType>>>,
     global_values: Rc<RefCell<Vec<Value>>>,
+    function_registry: Rc<RefCell<HashMap<u32, TypedNode>>>,
 }
 
 impl MacroEvaluator for RuntimeMacroEvaluator {
@@ -58,6 +80,8 @@ impl Environment {
             global_names: Rc::new(RefCell::new(HashMap::new())),
             global_types: Rc::new(RefCell::new(HashMap::new())),
             global_values: Rc::new(RefCell::new(Vec::new())),
+            function_registry: Rc::new(RefCell::new(HashMap::new())),
+            monomorph_cache: Rc::new(RefCell::new(HashMap::new())),
             debug_mode: false,
         };
         env.register_stdlib();
@@ -73,6 +97,7 @@ impl Environment {
             global_names: self.global_names.clone(),
             global_types: self.global_types.clone(),
             global_values: self.global_values.clone(),
+            function_registry: self.function_registry.clone(),
         };
         MacroExpander::new(MacroRegistry::new(), evaluator)
     }
@@ -136,7 +161,83 @@ impl Environment {
 
     /// Backend: Optimization (TCO, etc.)
     pub fn link(&self, node: TypedNode) -> TypedNode {
-        TCO::optimize(node)
+        // 1. Collect Lambdas (Populate the registry for the specializer)
+        LambdaCollector::collect(&node, &mut self.function_registry.borrow_mut());
+
+        // 2. Specialize
+        let specialized = self.specialize_node(node);
+        // let specialized = node;
+
+        // 3. Optimize
+        TCO::optimize(specialized)
+    }
+
+    fn specialize_node(&self, node: TypedNode) -> TypedNode {
+        let registry = Rc::new(EnvFunctionRegistry {
+            registry: self.function_registry.clone(),
+        });
+        
+        let rtl_lookup = Rc::new(|name: &str, args: &[StaticType]| intrinsics::lookup(name, args));
+        
+        // We need to construct a compiler callback that can recursively specialize and compile.
+        // To avoid complex self-capturing, we reconstruct the environment context needed.
+        let func_reg = self.function_registry.clone();
+        let mono_cache = self.monomorph_cache.clone();
+        let global_values = self.global_values.clone(); // Needed for VM/Closure creation
+        
+        let compiler = Rc::new(move |func_node: TypedNode, _arg_types: &[StaticType]| -> Result<(Value, StaticType), String> {
+            // 1. Specialize the body (Recursive)
+            // We recreate the specializer context here. 
+            // Note: This creates a new Specializer for each recursion, but they SHARE the 'mono_cache'.
+            let sub_registry = Rc::new(EnvFunctionRegistry { registry: func_reg.clone() });
+            let sub_rtl_lookup = Rc::new(|name: &str, args: &[StaticType]| intrinsics::lookup(name, args));
+            
+            // Note: We are passing 'None' as compiler to the inner specializer for now to prevent infinite recursion on cycles.
+            // A robust implementation would handle the recursion cycle or use a shared compiler reference.
+            // For 'tak', the recursion is handled by the cache or dynamic fallback.
+            let sub_specializer = Specializer::new(
+                Some(sub_registry),
+                None, // recursive compilation limit (depth 1) for safety
+                Some(sub_rtl_lookup),
+                Some(mono_cache.clone())
+            );
+            
+            let specialized_ast = sub_specializer.specialize(func_node);
+            
+            // 2. Optimize (TCO)
+            let optimized_ast = TCO::optimize(specialized_ast);
+            
+            // 3. Compile to Closure (VM)
+            // We run the VM once to evaluate the Lambda definition, producing a closure Value.
+            let mut vm = VM::new(global_values.clone());
+            
+            let compiled_val = match vm.run(&optimized_ast) {
+                Ok(v) => v,
+                Err(e) => return Err(format!("VM Error during specialization: {}", e)),
+            };
+
+            // We need the return type.
+            // For a Lambda, the type is stored in the node.
+            // But we need the return type of the FUNCTION (e.g. Int), not the type of the Lambda node (Method).
+            // Actually, Specializer expects (Value, ReturnType).
+            // If the specialized function returns Int, we return Int.
+            let ret_type = if let BoundKind::Lambda { body, .. } = &optimized_ast.kind {
+                body.ty.clone()
+            } else {
+                StaticType::Any
+            };
+
+            Ok((compiled_val, ret_type))
+        });
+        
+        let specializer = Specializer::new(
+            Some(registry),
+            Some(compiler),
+            Some(rtl_lookup),
+            Some(self.monomorph_cache.clone()),
+        );
+        
+        specializer.specialize(node)
     }
 
     /// Runtime: Execute the linked AST in the VM

src/ast/rtl/mod.rs

@@ -1,6 +1,7 @@
 pub mod core;
 pub mod datetime;
 pub mod type_registry;
+pub mod intrinsics;
 
 use crate::ast::environment::Environment;
 

src/bin/ast.rs

@@ -22,6 +22,10 @@ struct Cli {
     /// Update the benchmark baseline (Only allowed in Release mode)
     #[arg(short, long)]
     update_bench: bool,
+
+    /// Dump the compiled AST
+    #[arg(short, long)]
+    dump: bool,
 }
 
 fn main() {
@@ -45,10 +49,26 @@ fn main() {
     }
 
     if let Some(script_str) = cli.eval {
-        execute(&env, &script_str);
+        if cli.dump {
+            match env.dump_ast(&script_str) {
+                Ok(dump) => println!("{}", dump),
+                Err(e) => eprintln!("Error dumping AST: {}", e),
+            }
+        } else {
+            execute(&env, &script_str);
+        }
     } else if let Some(file_path) = cli.file {
         match fs::read_to_string(&file_path) {
-            Ok(content) => execute(&env, &content),
+            Ok(content) => {
+                if cli.dump {
+                    match env.dump_ast(&content) {
+                        Ok(dump) => println!("{}", dump),
+                        Err(e) => eprintln!("Error dumping AST: {}", e),
+                    }
+                } else {
+                    execute(&env, &content);
+                }
+            },
             Err(e) => {
                 eprintln!("Error reading file {:?}: {}", file_path, e);
                 std::process::exit(1);