Refactor Optimizer to use Folder

Extract common AST manipulation and folding logic into a new Folder
struct. This improves code organization and reusability. The Optimizer
now delegates these tasks to the Folder.
This commit is contained in:
Michael Schimmel
2026-02-25 19:57:00 +01:00
parent d3d1497c02
commit 7f64e7e6ea
4 changed files with 112 additions and 83 deletions
+761
View File
@@ -0,0 +1,761 @@
use crate::ast::compiler::bound_nodes::{
Address, AnalyzedNode, BoundKind, GlobalIdx, LocalSlot, UpvalueIdx,
};
use crate::ast::nodes::Node;
use crate::ast::types::{Purity, Value};
use crate::ast::vm::Closure;
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::rc::Rc;
use super::substitution_map::{SubstitutionMap, UsageInfo};
use super::folder::Folder;
pub struct Optimizer {
pub enabled: bool,
max_passes: usize,
pub globals: Option<Rc<RefCell<Vec<Value>>>>,
pub global_purity: Option<Rc<RefCell<HashMap<GlobalIdx, Purity>>>>,
pub lambda_registry: Option<Rc<RefCell<HashMap<GlobalIdx, AnalyzedNode>>>>,
}
struct PathTracker {
inlining_depth: usize,
inlining_stack: HashSet<GlobalIdx>,
identity_stack: HashSet<crate::ast::types::Identity>,
}
impl PathTracker {
fn new() -> Self {
Self {
inlining_depth: 0,
inlining_stack: HashSet::new(),
identity_stack: HashSet::new(),
}
}
fn enter_lambda(&mut self, identity: &crate::ast::types::Identity) -> bool {
if self.identity_stack.contains(identity) {
return false;
}
self.identity_stack.insert(identity.clone());
true
}
fn exit_lambda(&mut self, identity: &crate::ast::types::Identity) {
self.identity_stack.remove(identity);
}
}
impl Optimizer {
pub fn new(enabled: bool) -> Self {
Self {
enabled,
max_passes: 5,
globals: None,
global_purity: None,
lambda_registry: None,
}
}
pub fn with_globals(mut self, globals: Rc<RefCell<Vec<Value>>>) -> Self {
self.globals = Some(globals);
self
}
pub fn with_purity(mut self, purity: Rc<RefCell<HashMap<GlobalIdx, Purity>>>) -> Self {
self.global_purity = Some(purity);
self
}
pub fn with_registry(
mut self,
registry: Rc<RefCell<HashMap<GlobalIdx, AnalyzedNode>>>,
) -> Self {
self.lambda_registry = Some(registry);
self
}
pub fn optimize(&self, node: AnalyzedNode) -> AnalyzedNode {
if !self.enabled {
return node;
}
let mut current = node;
for _ in 0..self.max_passes {
let mut sub = SubstitutionMap::new();
let mut path = PathTracker::new();
let next = self.visit_node(current.clone(), &mut sub, &mut path);
if next == current {
break;
}
current = next;
}
current
}
fn visit_node(
&self,
node: AnalyzedNode,
sub: &mut SubstitutionMap,
path: &mut PathTracker,
) -> AnalyzedNode {
let folder = Folder::new(&self.globals);
let (new_kind, metrics) = match node.kind {
BoundKind::Get { addr, ref name } => {
if !sub.assigned.contains(&addr) {
// 1. Try inlining from current value substitution map (locals/globals/upvalues)
if let Some(val) = sub.get_value(&addr)
&& self.is_inlinable_value(val, addr)
{
return folder.make_constant_node(val.clone(), &node);
}
// 2. Try inlining from AST substitution map (pure expressions)
if let Some(inlined_node) = sub.ast_substitutions.get(&addr) {
return inlined_node.clone();
}
// 3. Fallback for Globals: check the actual VM environment
if let Address::Global(idx) = addr
&& let Some(globals_rc) = &self.globals
{
let globals = globals_rc.borrow();
if let Some(val) = globals.get(idx.0 as usize)
&& self.is_inlinable_value(val, addr)
{
return folder.make_constant_node(val.clone(), &node);
}
}
}
sub.used.insert(addr);
(
BoundKind::Get {
addr: sub.map_address(addr),
name: name.clone(),
},
node.ty.clone(),
)
}
BoundKind::Set { addr, value } => {
let value = Box::new(self.visit_node(*value, sub, path));
if let BoundKind::Constant(val) = &value.kind {
sub.add_value(addr, val.clone());
} else {
sub.remove_value(&addr);
}
(
BoundKind::Set {
addr: sub.map_address(addr),
value,
},
node.ty.clone(),
)
}
BoundKind::Define {
ref name,
addr,
kind,
ref value,
ref captured_by,
} => {
let value_opt = Box::new(self.visit_node((**value).clone(), sub, path));
if let Address::Local(slot) = addr
&& !captured_by.is_empty()
{
sub.captured_slots.insert(slot);
}
if let BoundKind::Constant(val) = &value_opt.kind {
sub.add_value(addr, val.clone());
} else {
sub.remove_value(&addr);
}
if let Address::Global(global_index) = addr
&& value_opt.ty.purity > Purity::Impure
&& let Some(purity_rc) = &self.global_purity
{
purity_rc
.borrow_mut()
.insert(global_index, value_opt.ty.purity);
}
(
BoundKind::Define {
name: name.clone(),
addr: sub.map_address(addr),
kind,
value: value_opt,
captured_by: captured_by.clone(),
},
node.ty.clone(),
)
}
BoundKind::Call { callee, args } => {
let callee = self.visit_node(*callee, sub, path);
let args = self.visit_node(*args, sub, path);
if self.enabled {
if let BoundKind::Lambda {
params,
body,
upvalues,
positional_count,
} = &callee.kind
&& upvalues.is_empty()
&& positional_count.is_some()
&& path.inlining_depth < 5
&& !callee.ty.is_recursive
&& path.enter_lambda(&callee.identity)
{
path.inlining_depth += 1;
let collapsed = self.try_beta_reduce_with_sub(
params,
&args,
(**body).clone(),
&mut SubstitutionMap::new(),
path,
);
path.inlining_depth -= 1;
path.exit_lambda(&callee.identity);
if let Some(res) = collapsed {
return res;
}
}
if let BoundKind::Get {
addr: Address::Global(idx),
..
} = &callee.kind
&& let Some(registry_rc) = &self.lambda_registry
&& path.inlining_depth < 5
&& !path.inlining_stack.contains(idx)
{
let registry = registry_rc.borrow();
if let Some(lambda_node) = registry.get(idx)
&& let BoundKind::Lambda {
params,
body,
upvalues,
positional_count,
} = &lambda_node.kind
&& upvalues.is_empty()
&& positional_count.is_some()
&& !lambda_node.ty.is_recursive
{
let mut inner_sub = SubstitutionMap::new();
path.inlining_stack.insert(*idx);
path.inlining_depth += 1;
let collapsed = self.try_beta_reduce_with_sub(
params.as_ref(),
&args,
body.as_ref().clone(),
&mut inner_sub,
path,
);
path.inlining_depth -= 1;
path.inlining_stack.remove(idx);
if let Some(res) = collapsed {
return res;
}
}
}
if let BoundKind::Constant(Value::Object(ref obj)) = callee.kind
&& path.inlining_depth < 5
&& let Some(closure) = obj.as_any().downcast_ref::<Closure>()
&& (closure.upvalues.is_empty()
|| closure.function_node.ty.purity >= Purity::SideEffectFree)
&& !closure.function_node.ty.is_recursive
&& path.enter_lambda(&closure.function_node.identity)
{
let mut closure_sub = SubstitutionMap::new();
for (i, cell) in closure.upvalues.iter().enumerate() {
closure_sub.add_value(
Address::Upvalue(UpvalueIdx(i as u32)),
cell.borrow().clone(),
);
}
path.inlining_depth += 1;
let inlined_body = self.visit_node(
(*closure.function_node).clone(),
&mut closure_sub,
path,
);
let collapsed = self.try_beta_reduce_with_sub(
&closure.parameter_node,
&args,
inlined_body,
&mut closure_sub,
path,
);
path.inlining_depth -= 1;
path.exit_lambda(&closure.function_node.identity);
if let Some(res) = collapsed {
return res;
}
}
if let Some(folded) = folder.try_fold_pure(&callee, &args) {
return folded;
}
}
(
BoundKind::Call {
callee: Box::new(callee),
args: Box::new(args),
},
node.ty.clone(),
)
}
BoundKind::Again { args } => {
let args = self.visit_node(*args, sub, path);
(
BoundKind::Again {
args: Box::new(args),
},
node.ty.clone(),
)
}
BoundKind::If {
ref cond,
ref then_br,
ref else_br,
} => {
let cond_opt = self.visit_node((**cond).clone(), sub, path);
if self.enabled
&& let BoundKind::Constant(ref val) = cond_opt.kind
{
if val.is_truthy() {
return self.visit_node((**then_br).clone(), sub, path);
} else if let Some(else_node) = else_br {
return self.visit_node((**else_node).clone(), sub, path);
} else {
return folder.make_nop_node(&node);
}
}
let then_br = Box::new(self.visit_node((**then_br).clone(), sub, path));
let else_br = else_br
.as_ref()
.map(|e| Box::new(self.visit_node((**e).clone(), sub, path)));
(
BoundKind::If {
cond: Box::new(cond_opt),
then_br,
else_br,
},
node.ty.clone(),
)
}
BoundKind::Block { ref exprs } => {
let mut info = UsageInfo::default();
if !exprs.is_empty() {
for e in exprs {
info.collect(e);
}
}
sub.assigned.extend(info.assigned.iter().cloned());
let mut new_exprs = Vec::with_capacity(exprs.len());
let last_idx = exprs.len().saturating_sub(1);
for (i, e) in exprs.iter().enumerate() {
let is_last = i == last_idx;
if self.enabled && !is_last {
let removable = match &e.kind {
BoundKind::Define { addr, value, .. } => {
!info.is_used(addr)
&& (if let Address::Local(slot) = addr {
!sub.captured_slots.contains(slot)
} else {
true
})
&& (value.ty.purity >= Purity::SideEffectFree
|| matches!(value.kind, BoundKind::Lambda { .. }))
}
BoundKind::Set { addr, value, .. } => {
!info.is_used(addr)
&& (if let Address::Local(slot) = addr {
!sub.captured_slots.contains(slot)
} else {
true
})
&& value.ty.purity >= Purity::SideEffectFree
}
_ => e.ty.purity >= Purity::SideEffectFree,
};
if removable {
continue;
}
}
let opt = self.visit_node(e.clone(), sub, path);
if self.enabled && matches!(opt.kind, BoundKind::Nop) && !is_last {
continue;
}
new_exprs.push(opt);
}
if self.enabled {
if new_exprs.is_empty() {
return folder.make_nop_node(&node);
} else if new_exprs.len() == 1 {
return new_exprs.pop().unwrap();
}
}
(BoundKind::Block { exprs: new_exprs }, node.ty.clone())
}
BoundKind::Lambda { .. } => {
let (params, original_upvalues, body, positional_count) =
if let BoundKind::Lambda {
params,
upvalues,
body,
positional_count,
} = &node.kind
{
(params, upvalues, body, positional_count)
} else {
unreachable!()
};
let mut info = UsageInfo::default();
info.collect(&node);
let mut new_upvalues = Vec::new();
let mut mapping = Vec::new();
let mut next_inner_subs = SubstitutionMap::new();
next_inner_subs.assigned = info.assigned;
for (old_idx, capture_addr) in original_upvalues.iter().enumerate() {
let mut inlined_val = None;
if let Some(val) = sub.get_value(capture_addr) {
inlined_val = Some(val.clone());
}
if let Address::Local(slot) = capture_addr {
sub.captured_slots.insert(*slot);
}
if let Some(val) = inlined_val {
next_inner_subs
.add_value(Address::Upvalue(UpvalueIdx(old_idx as u32)), val);
mapping.push(None);
} else {
mapping.push(Some(new_upvalues.len() as u32));
new_upvalues.push(sub.map_address(*capture_addr));
}
}
let params_node =
self.visit_node(params.as_ref().clone(), &mut next_inner_subs, path);
self.collect_parameter_slots(&params_node, &mut next_inner_subs);
let body_node = self.visit_node(body.as_ref().clone(), &mut next_inner_subs, path);
let reindexed_body = if new_upvalues.len() != original_upvalues.len() {
sub.reindex_upvalues(body_node, &mapping)
} else {
body_node
};
(
BoundKind::Lambda {
params: Rc::new(params_node),
upvalues: new_upvalues,
body: Rc::new(reindexed_body),
positional_count: *positional_count,
},
node.ty.clone(),
)
}
BoundKind::Destructure {
ref pattern,
ref value,
} => {
let val_opt = Box::new(self.visit_node((**value).clone(), sub, path));
let pat_opt = Box::new(self.visit_node((**pattern).clone(), sub, path));
let mut info = UsageInfo::default();
info.collect_pattern(&pat_opt);
for addr in info.assigned {
sub.remove_value(&addr);
}
(
BoundKind::Destructure {
pattern: pat_opt,
value: val_opt,
},
node.ty.clone(),
)
}
BoundKind::Tuple { elements } => {
let elements = elements
.into_iter()
.map(|e| self.visit_node(e, sub, path))
.collect();
(BoundKind::Tuple { elements }, node.ty.clone())
}
BoundKind::Record { fields } => {
let fields = fields
.into_iter()
.map(|(k, v)| (self.visit_node(k, sub, path), self.visit_node(v, sub, path)))
.collect();
(BoundKind::Record { fields }, node.ty.clone())
}
BoundKind::Expansion {
ref original_call,
ref bound_expanded,
} => {
path.inlining_depth += 1;
let bound_expanded =
Box::new(self.visit_node((**bound_expanded).clone(), sub, path));
path.inlining_depth -= 1;
(
BoundKind::Expansion {
original_call: original_call.clone(),
bound_expanded,
},
node.ty.clone(),
)
}
k => (k, node.ty.clone()),
};
Node {
identity: node.identity,
kind: new_kind,
ty: metrics,
}
}
fn collect_parameter_slots(&self, node: &AnalyzedNode, sub: &mut SubstitutionMap) {
match &node.kind {
BoundKind::Define {
addr: Address::Local(slot),
..
} => {
sub.slot_mapping.insert(*slot, *slot);
sub.next_slot = sub.next_slot.max(slot.0 + 1);
}
BoundKind::Tuple { elements } => {
for el in elements {
self.collect_parameter_slots(el, sub);
}
}
_ => {}
}
}
fn is_inlinable_value(&self, val: &Value, addr: Address) -> bool {
// 1. Basic check: is the type itself inlinable (Constants, pure Closures)?
let type_ok = match val {
Value::Int(_)
| Value::Float(_)
| Value::Bool(_)
| Value::Text(_)
| Value::Keyword(_)
| Value::DateTime(_) => true,
Value::Object(obj) => {
if let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
closure.upvalues.is_empty() && !closure.function_node.ty.is_recursive
} else {
false
}
}
_ => false,
};
if !type_ok {
return false;
}
// 2. For Globals: check if they are actually pure/stable
if let Address::Global(idx) = addr {
if let Some(purity_rc) = &self.global_purity {
let purity = purity_rc
.borrow()
.get(&idx)
.cloned()
.unwrap_or(Purity::Impure);
return purity >= Purity::Pure;
}
// If no purity info is available, we assume the global is unstable (safer)
return false;
}
true // Locals/Upvalues that reach here and passed the type check are inlinable
}
fn try_beta_reduce_with_sub(
&self,
params: &AnalyzedNode,
args: &AnalyzedNode,
body: AnalyzedNode,
sub: &mut SubstitutionMap,
path: &mut PathTracker,
) -> Option<AnalyzedNode> {
let mut body_usage = UsageInfo::default();
body_usage.collect(&body);
let mut arg_vals = Vec::new();
self.flatten_tuple(args.clone(), &mut arg_vals);
let mut slot_index = 0;
self.map_params_to_args(params, &arg_vals, &mut slot_index, sub, &body_usage);
if slot_index != arg_vals.len() {
return None;
}
// Verify that all used/assigned parameters have been substituted.
let mut param_slots = HashSet::new();
self.collect_parameter_slots_set(params, &mut param_slots);
for slot in param_slots {
let addr = Address::Local(slot);
if (body_usage.is_used(&addr) || body_usage.is_assigned(&addr))
&& sub.get_value(&addr).is_none()
&& !sub.ast_substitutions.contains_key(&addr)
{
return None;
}
}
if sub.values.is_empty() && sub.ast_substitutions.is_empty() && !arg_vals.is_empty() {
return None;
}
Some(self.visit_node(body, sub, path))
}
fn collect_parameter_slots_set(&self, node: &AnalyzedNode, slots: &mut HashSet<LocalSlot>) {
match &node.kind {
BoundKind::Define {
addr: Address::Local(slot),
..
} => {
slots.insert(*slot);
}
BoundKind::Tuple { elements } => {
for el in elements {
self.collect_parameter_slots_set(el, slots);
}
}
_ => {}
}
}
fn flatten_tuple(&self, node: AnalyzedNode, into: &mut Vec<AnalyzedNode>) {
match node.kind {
BoundKind::Tuple { elements } => {
for el in elements {
self.flatten_tuple(el, into);
}
}
BoundKind::Record { fields } => {
for (_, v) in fields {
self.flatten_tuple(v, into);
}
}
BoundKind::Nop => {}
_ => into.push(node),
}
}
fn map_params_to_args(
&self,
pattern: &AnalyzedNode,
args: &[AnalyzedNode],
offset: &mut usize,
sub: &mut SubstitutionMap,
body_usage: &UsageInfo,
) {
match &pattern.kind {
BoundKind::Define { addr, .. } => {
if let Some(arg) = args.get(*offset)
&& !body_usage.is_assigned(addr)
{
if let BoundKind::Constant(val) = &arg.kind {
sub.add_value(*addr, val.clone());
} else if let BoundKind::Lambda { upvalues, .. } = &arg.kind
&& upvalues.is_empty()
{
sub.add_ast_substitution(*addr, arg.clone());
} else if let BoundKind::Get {
addr: Address::Global(_),
..
} = &arg.kind
{
sub.add_ast_substitution(*addr, arg.clone());
}
}
// Ensure local slots are mapped even if no constant/AST is substituted
if let Address::Local(slot) = addr {
sub.slot_mapping.insert(*slot, *slot);
sub.next_slot = sub.next_slot.max(slot.0 + 1);
}
*offset += 1;
}
BoundKind::Tuple { elements } => {
// RECURSIVE DESTRUCTURING SUPPORT
// Check if the current argument at 'offset' is itself a Tuple or Record literal.
if let Some(arg) = args.get(*offset) {
let mut sub_args = Vec::new();
let is_compound = match &arg.kind {
BoundKind::Tuple { .. } | BoundKind::Record { .. } => {
self.flatten_tuple(arg.clone(), &mut sub_args);
true
}
_ => false,
};
if is_compound {
// Match inner elements against the flattened compound argument
let mut sub_offset = 0;
for el in elements {
self.map_params_to_args(
el,
&sub_args,
&mut sub_offset,
sub,
body_usage,
);
}
*offset += 1;
return;
}
}
// Fallback: Continue flat matching (original behavior)
for el in elements {
self.map_params_to_args(el, args, offset, sub, body_usage);
}
}
_ => {}
}
}
}