Brummel/RustAst
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Implement closure cracking and inlining

Browse Source 
Introduces a new optimizer pass that can "crack" closures, allowing for
more aggressive specialization. It also enables inlining of upvalues
that point to immutable global variables. This removes overhead for
higher-order functions and currying when arguments are statically
resolvable.
This commit is contained in:
Michael Schimmel
2026-02-21 18:49:55 +01:00
parent 74ea38248e
commit 0bbe35eeec
11 changed files with 386 additions and 286 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ast/compiler/specializer.rs
+4 -271
View File
@@ -90,11 +90,10 @@ impl Specializer {
let elements = elements.into_iter().map(|e| self.visit_node(e)).collect();
(BoundKind::Tuple { elements }, node.ty)
},
BoundKind::Record { fields } => {
let fields = fields.into_iter().map(|(k, v)| (self.visit_node(k), self.visit_node(v))).collect();
(BoundKind::Record { fields }, node.ty)
}
,
BoundKind::Record { fields } => {
let fields = fields.into_iter().map(|(k, v)| (self.visit_node(k), self.visit_node(v))).collect();
(BoundKind::Record { fields }, node.ty)
},
BoundKind::Expansion { original_call, bound_expanded } => {
let bound_expanded = Box::new(self.visit_node(*bound_expanded));
(BoundKind::Expansion { original_call, bound_expanded }, node.ty)
@@ -177,11 +176,9 @@ impl Specializer {
// Check constraints (no closures with state)
if let BoundKind::Lambda { upvalues, .. } = &func_node.kind {
if !upvalues.is_empty() {
// Cannot specialize stateful closures trivially
return (new_callee, new_args, original_ty);
}
} else {
// Not a lambda?
return (new_callee, new_args, original_ty);
}
@@ -196,7 +193,6 @@ impl Specializer {
self.cache.borrow_mut().insert(key, (res_val.clone(), res_ty.clone()));
// PERFORMANCE: Flatten the argument tuple to match the specialized signature.
// Since we are specializing, we can convert [[1 2] 3] into a flat [1 2 3] Tuple node.
let flat_elements = self.flatten_tuple(new_args.clone());
let flat_types = flat_elements.iter().map(|e| e.ty.clone()).collect();
let flattened_args = Node {
@@ -239,266 +235,3 @@ impl Specializer {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ast::types::{Identity, NodeIdentity, SourceLocation, StaticType, Value, Signature};
use crate::ast::compiler::bound_nodes::{BoundKind, Address, TypedNode};
use crate::ast::nodes::Symbol;
use std::rc::Rc;
fn make_identity() -> Identity {
Rc::new(NodeIdentity { location: SourceLocation { line: 0, col: 0 } })
}
fn make_typed_node(kind: BoundKind<StaticType>, ty: StaticType) -> TypedNode {
crate::ast::nodes::Node {
identity: make_identity(),
kind,
ty,
}
}
// Mock Registry
struct MockRegistry {
functions: HashMap<Address, BoundNode>,
}
impl MockRegistry {
fn new() -> Self {
Self { functions: HashMap::new() }
}
fn register(&mut self, addr: Address, node: BoundNode) {
self.functions.insert(addr, node);
}
}
impl FunctionRegistry for MockRegistry {
fn resolve(&self, addr: Address) -> Option<BoundNode> {
self.functions.get(&addr).cloned()
}
}
#[test]
fn test_specialize_compiles_user_function() {
// Setup Registry with a function definition
let mut registry = MockRegistry::new();
let addr = Address::Local(0);
let name = Symbol::from("test_func");
// Def: (fn [x] x) -- generic identity
let func_node = BoundNode {
identity: make_identity(),
kind: BoundKind::Lambda {
params: Rc::new(BoundNode {
identity: make_identity(),
kind: BoundKind::Tuple {
elements: vec![
BoundNode {
identity: make_identity(),
kind: BoundKind::Parameter { name: name.clone(), slot: 0 },
ty: ()
}
]
},
ty: ()
}),
upvalues: vec![],
body: Rc::new(BoundNode { identity: make_identity(), kind: BoundKind::Nop, ty: () }),
positional_count: Some(1),
},
ty: ()
};
registry.register(addr, func_node);
// Setup Compiler Mock
let compiler: CompileFunc = Rc::new(|_node: BoundNode, _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, None);
// Call(Get(Local(0)), Tuple([Arg(Int)]))
let callee = make_typed_node(BoundKind::Get { addr, name: name.clone() }, StaticType::Any);
let arg = make_typed_node(BoundKind::Constant(Value::Int(1)), StaticType::Int);
let args_tuple = make_typed_node(BoundKind::Tuple { elements: vec![arg] }, StaticType::Tuple(vec![StaticType::Int]));
let call_node = make_typed_node(
BoundKind::Call { callee: Box::new(callee), args: Box::new(args_tuple) },
StaticType::Any
);
let result = spec.specialize(call_node);
// Should be Call(Constant(12345), ...)
if let BoundKind::Call { callee, .. } = result.kind {
if let BoundKind::Constant(val) = callee.kind {
match val {
Value::Int(12345) => (),
_ => panic!("Expected compiled value 12345"),
}
} else {
panic!("Expected Constant callee");
}
} else {
panic!("Expected Call node");
}
}
#[test]
fn test_specialize_skips_unknown_types() {
let spec = Specializer::new(None, None, None, None);
let name0 = Symbol::from("f");
let name1 = Symbol::from("x");
// Call(Get(Local(0)), Tuple([Get(Local(1))])) where arg is Any
let callee = make_typed_node(BoundKind::Get { addr: Address::Local(0), name: name0 }, StaticType::Function(Box::new(Signature { params: StaticType::Tuple(vec![StaticType::Any]), ret: StaticType::Void })));
let arg = make_typed_node(BoundKind::Get { addr: Address::Local(1), name: name1 }, StaticType::Any);
let args_tuple = make_typed_node(BoundKind::Tuple { elements: vec![arg] }, StaticType::Tuple(vec![StaticType::Any]));
let call_node = make_typed_node(
BoundKind::Call { callee: Box::new(callee), args: Box::new(args_tuple) },
StaticType::Void
);
let result = spec.specialize(call_node);
// Should remain a generic Call because arg type is Any
if let BoundKind::Call { callee, .. } = result.kind {
if let BoundKind::Get { .. } = callee.kind {
// Correct: Still a Get, not a Constant(Function)
} else {
panic!("Expected generic Call to Get, got {:?}", callee.kind);
}
} else {
panic!("Expected Call node");
}
}
#[test]
fn test_specialize_uses_cache() {
// Setup cache with a pre-specialized function for (Int) -> Int
let spec = Specializer::new(None, None, None, None);
let addr = Address::Local(0);
let name = Symbol::from("cached_func");
let arg_types = vec![StaticType::Int];
let key = MonoCacheKey { address: addr, arg_types: arg_types.clone() };
// Mock a specialized function pointer
let specialized_val = Value::Int(999); // Dummy value representing function
let ret_ty = StaticType::Int;
spec.cache.borrow_mut().insert(key, (specialized_val.clone(), ret_ty.clone()));
// Create the call node: Call(Get(0), Tuple([Arg(Int)]))
let callee = make_typed_node(BoundKind::Get { addr, name }, StaticType::Any);
let arg = make_typed_node(BoundKind::Constant(Value::Int(1)), StaticType::Int);
let args_tuple = make_typed_node(BoundKind::Tuple { elements: vec![arg] }, StaticType::Tuple(vec![StaticType::Int]));
let call_node = make_typed_node(
BoundKind::Call { callee: Box::new(callee), args: Box::new(args_tuple) },
StaticType::Any
);
let result = spec.specialize(call_node);
// Should now be Call(Constant(999), ...)
if let BoundKind::Call { callee, .. } = result.kind {
if let BoundKind::Constant(val) = callee.kind {
match val {
Value::Int(999) => (), // Success
_ => panic!("Expected specialized value 999"),
}
} else {
panic!("Expected Constant callee, got {:?}", callee.kind);
}
} else {
panic!("Expected Call node");
}
}
#[test]
fn test_specialize_uses_rtl_lookup() {
// Setup RTL Lookup Mock
let rtl_lookup: RtlLookupFunc = Rc::new(|name, _args| {
if name == "rtl_func" {
Some((Value::Int(888), StaticType::Int))
} else {
None
}
});
let spec = Specializer::new(None, None, Some(rtl_lookup), None);
let addr = Address::Global(10);
let name = Symbol::from("rtl_func");
let callee = make_typed_node(BoundKind::Get { addr, name }, StaticType::Any);
let arg = make_typed_node(BoundKind::Constant(Value::Int(1)), StaticType::Int);
let args_tuple = make_typed_node(BoundKind::Tuple { elements: vec![arg] }, StaticType::Tuple(vec![StaticType::Int]));
let call_node = make_typed_node(
BoundKind::Call { callee: Box::new(callee), args: Box::new(args_tuple) },
StaticType::Any
);
let result = spec.specialize(call_node);
if let BoundKind::Call { callee, .. } = result.kind {
if let BoundKind::Constant(val) = callee.kind {
match val {
Value::Int(888) => (),
_ => panic!("Expected RTL value 888"),
}
} else {
panic!("Expected Constant callee from RTL");
}
} else {
panic!("Expected Call node");
}
}
#[test]
fn test_specialize_preserves_tail_call() {
let spec = Specializer::new(None, None, None, None);
let addr = Address::Local(0);
let name = Symbol::from("tail_func");
let arg_types = vec![StaticType::Int];
let key = MonoCacheKey { address: addr, arg_types: arg_types.clone() };
let specialized_val = Value::Int(777);
let ret_ty = StaticType::Int;
spec.cache.borrow_mut().insert(key, (specialized_val.clone(), ret_ty.clone()));
let callee = make_typed_node(BoundKind::Get { addr, name }, StaticType::Any);
let arg = make_typed_node(BoundKind::Constant(Value::Int(1)), StaticType::Int);
let args_tuple = make_typed_node(BoundKind::Tuple { elements: vec![arg] }, StaticType::Tuple(vec![StaticType::Int]));
// Use TailCall here
let call_node = make_typed_node(
BoundKind::TailCall { callee: Box::new(callee), args: Box::new(args_tuple) },
StaticType::Any
);
let result = spec.specialize(call_node);
if let BoundKind::TailCall { callee, .. } = result.kind {
if let BoundKind::Constant(val) = callee.kind {
match val {
Value::Int(777) => (),
_ => panic!("Expected specialized value 777"),
}
} else {
panic!("Expected Constant callee");
}
} else {
panic!("Expected TailCall node, got {:?}", result.kind);
}
}
}