From 1b49719d31bd283930fa801115cd50c2ce832590 Mon Sep 17 00:00:00 2001 From: Michael Schimmel Date: Thu, 19 Feb 2026 22:26:24 +0100 Subject: [PATCH] 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>`. 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 | 1 + src/ast/compiler/specializer.rs | 41 +++++++------ src/ast/environment.rs | 103 +++++++++++++++++++++++++++++++- src/ast/rtl/mod.rs | 1 + src/bin/ast.rs | 24 +++++++- 5 files changed, 148 insertions(+), 22 deletions(-) diff --git a/src/ast/compiler/mod.rs b/src/ast/compiler/mod.rs index b362d28..ae283f4 100644 --- a/src/ast/compiler/mod.rs +++ b/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::*; diff --git a/src/ast/compiler/specializer.rs b/src/ast/compiler/specializer.rs index ce48485..3c7f4b5 100644 --- a/src/ast/compiler/specializer.rs +++ b/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, } -pub type CompileFunc = Rc, &[StaticType]) -> Result<(Value, StaticType), String>>; +pub type CompileFunc = Rc Result<(Value, StaticType), String>>; pub type RtlLookupFunc = Rc Option<(Value, StaticType)>>; pub trait FunctionRegistry { - fn resolve(&self, addr: Address) -> Option>; + fn resolve(&self, addr: Address) -> Option; } +pub type MonoCache = HashMap; + pub struct Specializer { - cache: RefCell>, + pub cache: Rc>, registry: Option>, compiler: Option, rtl_lookup: Option, @@ -30,9 +32,10 @@ impl Specializer { registry: Option>, compiler: Option, rtl_lookup: Option, + cache: Option>>, ) -> 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>, + functions: HashMap, } 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> { + fn resolve(&self, addr: Address) -> Option { 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, _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"); diff --git a/src/ast/environment.rs b/src/ast/environment.rs index c39b57d..44dbd73 100644 --- a/src/ast/environment.rs +++ b/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>>, pub global_types: Rc>>, pub global_values: Rc>>, + pub function_registry: Rc>>, + pub monomorph_cache: Rc>, pub debug_mode: bool, } +struct EnvFunctionRegistry { + registry: Rc>>, +} + +impl FunctionRegistry for EnvFunctionRegistry { + fn resolve(&self, addr: Address) -> Option { + 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>>, global_types: Rc>>, global_values: Rc>>, + function_registry: Rc>>, } 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 diff --git a/src/ast/rtl/mod.rs b/src/ast/rtl/mod.rs index 3617682..900520a 100644 --- a/src/ast/rtl/mod.rs +++ b/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; diff --git a/src/bin/ast.rs b/src/bin/ast.rs index 7eff9e9..2786a16 100644 --- a/src/bin/ast.rs +++ b/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);