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.
This commit is contained in:
Michael Schimmel
2026-02-19 22:26:24 +01:00
parent 55502cbb69
commit 1b49719d31
5 changed files with 148 additions and 22 deletions
+1
View File
@@ -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::*;
+22 -19
View File
@@ -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");