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
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user