feat: Implement closure cracking and inlining
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
@@ -23,10 +23,9 @@ impl<'a> LambdaCollector<'a> {
|
||||
}
|
||||
|
||||
BoundKind::DefGlobal { global_index, value, .. } => {
|
||||
// If we define a global that is a lambda, register its BODY as the template.
|
||||
// This allows the VM to skip the Lambda/Parameter nodes during execution.
|
||||
if let BoundKind::Lambda { body, .. } = &value.kind {
|
||||
self.registry.insert(*global_index, (**body).clone());
|
||||
// Register the full Lambda node as the template.
|
||||
if let BoundKind::Lambda { .. } = &value.kind {
|
||||
self.registry.insert(*global_index, (**value).clone());
|
||||
}
|
||||
self.visit(value);
|
||||
}
|
||||
@@ -34,9 +33,9 @@ impl<'a> LambdaCollector<'a> {
|
||||
BoundKind::Set { addr, value } => {
|
||||
// Also track assignments to globals if they hold lambdas.
|
||||
if let Address::Global(global_index) = addr
|
||||
&& let BoundKind::Lambda { body, .. } = &value.kind
|
||||
&& let BoundKind::Lambda { .. } = &value.kind
|
||||
{
|
||||
self.registry.insert(*global_index, (**body).clone());
|
||||
self.registry.insert(*global_index, (**value).clone());
|
||||
}
|
||||
self.visit(value);
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ pub mod dumper;
|
||||
pub mod macros;
|
||||
pub mod type_checker;
|
||||
pub mod specializer;
|
||||
pub mod optimizer;
|
||||
pub mod lambda_collector;
|
||||
|
||||
pub use binder::*;
|
||||
@@ -16,3 +17,4 @@ pub use dumper::*;
|
||||
pub use macros::*;
|
||||
pub use type_checker::*;
|
||||
pub use specializer::*;
|
||||
pub use optimizer::*;
|
||||
|
||||
@@ -0,0 +1,274 @@
|
||||
use std::rc::Rc;
|
||||
use crate::ast::compiler::bound_nodes::{BoundKind, TypedNode, Address};
|
||||
use crate::ast::types::{Value, StaticType};
|
||||
use crate::ast::vm::Closure;
|
||||
use crate::ast::nodes::Node;
|
||||
|
||||
/// The Optimizer performs Phase 2: Partial Evaluation & Closure Cracking.
|
||||
pub struct Optimizer {
|
||||
/// 0: None, 1: Cracking (Stateless transformation), 2: Aggressive (Collapsing)
|
||||
pub level: u32,
|
||||
max_passes: usize,
|
||||
}
|
||||
|
||||
impl Optimizer {
|
||||
pub fn new(level: u32) -> Self {
|
||||
Self { level, max_passes: 3 }
|
||||
}
|
||||
|
||||
pub fn optimize(&self, node: TypedNode) -> TypedNode {
|
||||
if self.level == 0 {
|
||||
return node;
|
||||
}
|
||||
|
||||
let mut current = node;
|
||||
for _ in 0..self.max_passes {
|
||||
let next = self.visit_node(current.clone());
|
||||
current = next;
|
||||
}
|
||||
current
|
||||
}
|
||||
|
||||
fn visit_node(&self, node: TypedNode) -> TypedNode {
|
||||
let (new_kind, new_ty) = match node.kind {
|
||||
BoundKind::Call { callee, args } => {
|
||||
let callee = self.visit_node(*callee);
|
||||
let args = self.visit_node(*args);
|
||||
|
||||
// --- Cracking & Collapsing ---
|
||||
if let BoundKind::Constant(Value::Object(ref obj)) = callee.kind {
|
||||
if let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
|
||||
// Level 1+: Inlining Upvalues (Stateless transformation)
|
||||
// Structure is PRESERVED (Call still exists), but code is specialized.
|
||||
let body = (*closure.function_node).clone();
|
||||
let inlined_body = self.inline_upvalues(body, &closure.upvalues, 0);
|
||||
|
||||
// Level 2: Aggressive Collapsing
|
||||
if self.level >= 2 {
|
||||
// 1. Collapse 0-parameter calls
|
||||
if let Some(0) = closure.positional_count {
|
||||
let is_empty_args = match &args.kind {
|
||||
BoundKind::Nop => true,
|
||||
BoundKind::Tuple { elements } => elements.is_empty(),
|
||||
_ => false,
|
||||
};
|
||||
if is_empty_args {
|
||||
return self.visit_node(inlined_body);
|
||||
}
|
||||
}
|
||||
// 2. Beta-Reduction (TODO: Inlining parameters)
|
||||
}
|
||||
|
||||
// Stateless transformation (Cracking)
|
||||
let cracked_lambda = Node {
|
||||
identity: callee.identity.clone(),
|
||||
ty: callee.ty.clone(),
|
||||
kind: BoundKind::Lambda {
|
||||
params: closure.parameter_node.clone(),
|
||||
upvalues: vec![],
|
||||
body: Rc::new(inlined_body),
|
||||
positional_count: closure.positional_count,
|
||||
},
|
||||
};
|
||||
|
||||
return Node {
|
||||
identity: node.identity,
|
||||
kind: BoundKind::Call {
|
||||
callee: Box::new(cracked_lambda),
|
||||
args: Box::new(args)
|
||||
},
|
||||
ty: node.ty,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Level 2: Intrinsic Folding (Arithmetic on constants)
|
||||
if self.level >= 2 {
|
||||
if let Some(folded) = self.try_fold_intrinsic(&callee, &args) {
|
||||
return folded;
|
||||
}
|
||||
}
|
||||
|
||||
(BoundKind::Call { callee: Box::new(callee), args: Box::new(args) }, node.ty)
|
||||
},
|
||||
|
||||
BoundKind::If { cond, then_br, else_br } => {
|
||||
let cond = self.visit_node(*cond);
|
||||
if self.level >= 2 {
|
||||
if let BoundKind::Constant(ref val) = cond.kind {
|
||||
if val.is_truthy() {
|
||||
return self.visit_node(*then_br);
|
||||
} else if let Some(else_node) = else_br {
|
||||
return self.visit_node(*else_node);
|
||||
} else {
|
||||
return Node {
|
||||
identity: node.identity,
|
||||
kind: BoundKind::Nop,
|
||||
ty: StaticType::Void,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let then_br = Box::new(self.visit_node(*then_br));
|
||||
let else_br = else_br.map(|e| Box::new(self.visit_node(*e)));
|
||||
(BoundKind::If { cond: Box::new(cond), then_br, else_br }, node.ty)
|
||||
},
|
||||
|
||||
BoundKind::Block { exprs } => {
|
||||
let mut new_exprs = Vec::with_capacity(exprs.len());
|
||||
for e in exprs {
|
||||
let opt = self.visit_node(e);
|
||||
if self.level >= 2 && matches!(opt.kind, BoundKind::Nop) {
|
||||
continue;
|
||||
}
|
||||
new_exprs.push(opt);
|
||||
}
|
||||
|
||||
if self.level >= 2 {
|
||||
if new_exprs.is_empty() {
|
||||
return Node { identity: node.identity, kind: BoundKind::Nop, ty: StaticType::Void };
|
||||
} else if new_exprs.len() == 1 {
|
||||
return new_exprs.pop().unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
let ty = if new_exprs.is_empty() { StaticType::Void } else { new_exprs.last().unwrap().ty.clone() };
|
||||
(BoundKind::Block { exprs: new_exprs }, ty)
|
||||
},
|
||||
|
||||
// Recursive Traversal
|
||||
BoundKind::Lambda { params, upvalues, body, positional_count } => {
|
||||
let params = Rc::new(self.visit_node(params.as_ref().clone()));
|
||||
let body = Rc::new(self.visit_node(body.as_ref().clone()));
|
||||
(BoundKind::Lambda { params, upvalues, body, positional_count }, node.ty)
|
||||
},
|
||||
BoundKind::DefLocal { name, slot, value, captured_by } => {
|
||||
let value = Box::new(self.visit_node(*value));
|
||||
(BoundKind::DefLocal { name, slot, value, captured_by }, node.ty)
|
||||
},
|
||||
BoundKind::DefGlobal { name, global_index, value } => {
|
||||
let value = Box::new(self.visit_node(*value));
|
||||
(BoundKind::DefGlobal { name, global_index, value }, node.ty)
|
||||
},
|
||||
BoundKind::Set { addr, value } => {
|
||||
let value = Box::new(self.visit_node(*value));
|
||||
(BoundKind::Set { addr, value }, node.ty)
|
||||
},
|
||||
BoundKind::Tuple { elements } => {
|
||||
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::Expansion { original_call, bound_expanded } => {
|
||||
let bound_expanded = Box::new(self.visit_node(*bound_expanded));
|
||||
(BoundKind::Expansion { original_call, bound_expanded }, node.ty)
|
||||
},
|
||||
|
||||
k => (k, node.ty),
|
||||
};
|
||||
|
||||
Node {
|
||||
identity: node.identity,
|
||||
kind: new_kind,
|
||||
ty: new_ty,
|
||||
}
|
||||
}
|
||||
|
||||
fn try_fold_intrinsic(&self, callee: &TypedNode, args: &TypedNode) -> Option<TypedNode> {
|
||||
// Simple constant folding for + and - on Ints
|
||||
if let BoundKind::Get { name, .. } = &callee.kind {
|
||||
if let BoundKind::Tuple { elements } = &args.kind {
|
||||
if elements.len() == 2 {
|
||||
if let (BoundKind::Constant(Value::Int(a)), BoundKind::Constant(Value::Int(b))) = (&elements[0].kind, &elements[1].kind) {
|
||||
let res = match &*name.name {
|
||||
"+" => Some(Value::Int(a + b)),
|
||||
"-" => Some(Value::Int(a - b)),
|
||||
"*" => Some(Value::Int(a * b)),
|
||||
"/" if *b != 0 => Some(Value::Int(a / b)),
|
||||
_ => None
|
||||
};
|
||||
if let Some(val) = res {
|
||||
return Some(Node {
|
||||
identity: callee.identity.clone(),
|
||||
ty: StaticType::Int,
|
||||
kind: BoundKind::Constant(val),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn inline_upvalues(&self, node: TypedNode, upvalues: &[Rc<std::cell::RefCell<Value>>], depth: usize) -> TypedNode {
|
||||
let (new_kind, new_ty) = match node.kind {
|
||||
BoundKind::Get { addr: Address::Upvalue(idx), name } => {
|
||||
if depth == 0 {
|
||||
if let Some(cell) = upvalues.get(idx as usize) {
|
||||
let val = cell.borrow().clone();
|
||||
let ty = val.static_type();
|
||||
(BoundKind::Constant(val), ty)
|
||||
} else {
|
||||
(BoundKind::Get { addr: Address::Upvalue(idx), name }, node.ty)
|
||||
}
|
||||
} else {
|
||||
(BoundKind::Get { addr: Address::Upvalue(idx), name }, node.ty)
|
||||
}
|
||||
},
|
||||
|
||||
BoundKind::If { cond, then_br, else_br } => {
|
||||
let cond = Box::new(self.inline_upvalues(*cond, upvalues, depth));
|
||||
let then_br = Box::new(self.inline_upvalues(*then_br, upvalues, depth));
|
||||
let else_br = else_br.map(|e| Box::new(self.inline_upvalues(*e, upvalues, depth)));
|
||||
(BoundKind::If { cond, then_br, else_br }, node.ty)
|
||||
},
|
||||
BoundKind::Block { exprs } => {
|
||||
let exprs = exprs.into_iter().map(|e| self.inline_upvalues(e, upvalues, depth)).collect();
|
||||
(BoundKind::Block { exprs }, node.ty)
|
||||
},
|
||||
BoundKind::Call { callee, args } => {
|
||||
let callee = Box::new(self.inline_upvalues(*callee, upvalues, depth));
|
||||
let args = Box::new(self.inline_upvalues(*args, upvalues, depth));
|
||||
(BoundKind::Call { callee, args }, node.ty)
|
||||
},
|
||||
BoundKind::TailCall { callee, args } => {
|
||||
let callee = Box::new(self.inline_upvalues(*callee, upvalues, depth));
|
||||
let args = Box::new(self.inline_upvalues(*args, upvalues, depth));
|
||||
(BoundKind::TailCall { callee, args }, node.ty)
|
||||
},
|
||||
BoundKind::DefLocal { name, slot, value, captured_by } => {
|
||||
let value = Box::new(self.inline_upvalues(*value, upvalues, depth));
|
||||
(BoundKind::DefLocal { name, slot, value, captured_by }, node.ty)
|
||||
},
|
||||
BoundKind::Set { addr, value } => {
|
||||
let value = Box::new(self.inline_upvalues(*value, upvalues, depth));
|
||||
(BoundKind::Set { addr, value }, node.ty)
|
||||
},
|
||||
BoundKind::Tuple { elements } => {
|
||||
let elements = elements.into_iter().map(|e| self.inline_upvalues(e, upvalues, depth)).collect();
|
||||
(BoundKind::Tuple { elements }, node.ty)
|
||||
},
|
||||
BoundKind::Record { fields } => {
|
||||
let fields = fields.into_iter().map(|(k, v)| (self.inline_upvalues(k, upvalues, depth), self.inline_upvalues(v, upvalues, depth))).collect();
|
||||
(BoundKind::Record { fields }, node.ty)
|
||||
},
|
||||
BoundKind::Lambda { params, upvalues: lambda_upvalues, body, positional_count } => {
|
||||
let body = Rc::new(self.inline_upvalues(body.as_ref().clone(), upvalues, depth + 1));
|
||||
(BoundKind::Lambda { params, upvalues: lambda_upvalues, body, positional_count }, node.ty)
|
||||
},
|
||||
|
||||
k => (k, node.ty),
|
||||
};
|
||||
|
||||
Node {
|
||||
identity: node.identity,
|
||||
kind: new_kind,
|
||||
ty: new_ty,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+21
-9
@@ -13,6 +13,7 @@ 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::compiler::optimizer::Optimizer;
|
||||
use crate::ast::rtl;
|
||||
use crate::ast::rtl::intrinsics;
|
||||
use crate::ast::compiler::bound_nodes::{Address, BoundNode};
|
||||
@@ -24,6 +25,7 @@ pub struct Environment {
|
||||
pub function_registry: Rc<RefCell<HashMap<u32, BoundNode>>>,
|
||||
pub monomorph_cache: Rc<RefCell<MonoCache>>,
|
||||
pub debug_mode: bool,
|
||||
pub optimization_level: u32,
|
||||
}
|
||||
|
||||
struct EnvFunctionRegistry {
|
||||
@@ -81,6 +83,7 @@ impl Environment {
|
||||
function_registry: Rc::new(RefCell::new(HashMap::new())),
|
||||
monomorph_cache: Rc::new(RefCell::new(HashMap::new())),
|
||||
debug_mode: false,
|
||||
optimization_level: 0,
|
||||
};
|
||||
env.register_stdlib();
|
||||
env
|
||||
@@ -161,11 +164,15 @@ impl Environment {
|
||||
|
||||
/// Backend: Optimization (TCO, etc.)
|
||||
pub fn link(&self, node: TypedNode) -> TypedNode {
|
||||
// 1. Specialize
|
||||
// 1. Specialize (Always performed for correctness)
|
||||
let specialized = self.specialize_node(node);
|
||||
|
||||
// 2. Optimize
|
||||
TCO::optimize(specialized)
|
||||
// 2. Optimize (Level 1: Cracking, Level 2: Collapsing)
|
||||
let optimizer = Optimizer::new(self.optimization_level);
|
||||
let optimized = optimizer.optimize(specialized);
|
||||
|
||||
// 3. TCO (Always performed)
|
||||
TCO::optimize(optimized)
|
||||
}
|
||||
|
||||
fn specialize_node(&self, node: TypedNode) -> TypedNode {
|
||||
@@ -179,6 +186,7 @@ impl Environment {
|
||||
let mono_cache = self.monomorph_cache.clone();
|
||||
let global_values = self.global_values.clone();
|
||||
let global_types = self.global_types.clone();
|
||||
let opt_level = self.optimization_level;
|
||||
|
||||
let compiler = Rc::new(move |func_template: BoundNode, arg_types: &[StaticType]| -> Result<(Value, StaticType), String> {
|
||||
// 1. Re-TypeCheck the template with concrete argument types
|
||||
@@ -198,18 +206,22 @@ impl Environment {
|
||||
|
||||
let specialized_ast = sub_specializer.specialize(retyped_ast);
|
||||
|
||||
// 3. Optimize (TCO)
|
||||
let optimized_ast = TCO::optimize(specialized_ast);
|
||||
// 3. Optimize (Phase 2: Cracking & Folding)
|
||||
let optimizer = Optimizer::new(opt_level);
|
||||
let optimized_ast = optimizer.optimize(specialized_ast);
|
||||
|
||||
// 4. Compile to Value (VM)
|
||||
// 4. TCO
|
||||
let tco_ast = TCO::optimize(optimized_ast);
|
||||
|
||||
// 5. Compile to Value (VM)
|
||||
let mut vm = VM::new(global_values.clone());
|
||||
let compiled_val = match vm.run(&optimized_ast) {
|
||||
let compiled_val = match vm.run(&tco_ast) {
|
||||
Ok(v) => v,
|
||||
Err(e) => return Err(format!("VM Error during specialization: {}", e)),
|
||||
};
|
||||
|
||||
// 5. Determine correct return type from the newly inferred function signature
|
||||
let ret_type = if let StaticType::Function(sig) = &optimized_ast.ty {
|
||||
// 6. Determine correct return type from the newly inferred function signature
|
||||
let ret_type = if let StaticType::Function(sig) = &tco_ast.ty {
|
||||
sig.ret.clone()
|
||||
} else {
|
||||
StaticType::Any
|
||||
|
||||
@@ -30,11 +30,16 @@ struct Cli {
|
||||
/// Run with TracingObserver enabled
|
||||
#[arg(short, long)]
|
||||
trace: bool,
|
||||
|
||||
/// Optimization level (0: None, 1: Cracking, 2: Aggressive)
|
||||
#[arg(short, long, default_value_t = 0)]
|
||||
opt: u32,
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let cli = Cli::parse();
|
||||
let mut env = Environment::new();
|
||||
env.optimization_level = cli.opt;
|
||||
|
||||
if cli.bench || cli.update_bench {
|
||||
if cfg!(debug_assertions) {
|
||||
|
||||
Reference in New Issue
Block a user