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:
@@ -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::*;
|
||||
|
||||
@@ -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");
|
||||
|
||||
Reference in New Issue
Block a user