use crate::ast::nodes::{ Address, AnalyzedNode, AssignBinding, GlobalAnalyzedRegistry, IdentifierBinding, LambdaBinding, Node, NodeKind, UpvalueIdx, }; use crate::ast::types::{Purity, Value}; use crate::ast::vm::Closure; use std::cell::RefCell; use std::rc::Rc; use super::folder::Folder; use super::inliner::Inliner; use super::substitution_map::SubstitutionMap; use super::utils::{collect_pattern_addrs, PathTracker, UsageInfo}; pub struct Optimizer { pub enabled: bool, max_passes: usize, pub globals: Option>>>, pub root_purity: Option>>>, pub lambda_registry: Option>>, } impl Optimizer { pub fn new(enabled: bool) -> Self { Self { enabled, max_passes: 5, globals: None, root_purity: None, lambda_registry: None, } } pub fn with_globals(mut self, globals: Rc>>) -> Self { self.globals = Some(globals); self } pub fn with_purity(mut self, purity: Rc>>) -> Self { self.root_purity = Some(purity); self } pub fn with_registry( mut self, registry: Rc>, ) -> Self { self.lambda_registry = Some(registry); self } pub fn optimize(&self, node: AnalyzedNode) -> AnalyzedNode { if !self.enabled { return node; } let mut current = Rc::new(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 Rc::ptr_eq(&next, ¤t) { break; } current = next; } // Unwrap the final Rc if we are at the end, or return a clone of the inner node. // Since AnalyzedNode is small now (header + Rcs), cloning is cheap. (*current).clone() } fn try_inline( &self, params: &AnalyzedNode, arg_nodes: &[Rc], body: &AnalyzedNode, sub: &mut SubstitutionMap, path: &mut PathTracker, base_sub: Option, ) -> Option> { let inliner = Inliner::new(&self.globals, &self.root_purity); let mut inner_sub = base_sub.unwrap_or_else(|| sub.new_for_inlining()); if inliner.prepare_beta_reduction(params, arg_nodes, body, &mut inner_sub).is_some() { let res = self.visit_node(Rc::new(body.clone()), &mut inner_sub, path); // Sync back state to parent substitution map sub.next_slot = inner_sub.next_slot; sub.used.extend(inner_sub.used.iter().cloned()); sub.assigned.extend(inner_sub.assigned.iter().cloned()); sub.captured_slots.extend(inner_sub.captured_slots.iter().cloned()); Some(res) } else { None } } fn visit_node( &self, node_rc: Rc, sub: &mut SubstitutionMap, path: &mut PathTracker, ) -> Rc { let node = &*node_rc; let folder = Folder::new(&self.globals); let inliner = Inliner::new(&self.globals, &self.root_purity); let (new_kind, metrics) = match &node.kind { NodeKind::Identifier { symbol, binding } => { let addr = match binding { IdentifierBinding::Reference(addr) => *addr, IdentifierBinding::Declaration { addr, .. } => *addr, }; if !sub.assigned.contains(&addr) { // 1. Try inlining from current value substitution map (locals/globals/upvalues) if let Some(val) = sub.get_value(&addr) && inliner.is_inlinable_value(val, addr) { return Rc::new(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) && inliner.is_inlinable_value(val, addr) { return Rc::new(folder.make_constant_node(val.clone(), node)); } } } sub.used.insert(addr); let new_binding = match binding { IdentifierBinding::Reference(_) => { IdentifierBinding::Reference(sub.map_address(addr)) } IdentifierBinding::Declaration { kind, .. } => { IdentifierBinding::Declaration { addr: sub.map_address(addr), kind: *kind, } } }; ( NodeKind::Identifier { symbol: symbol.clone(), binding: new_binding, }, node.ty.clone(), ) } NodeKind::FieldAccessor(k) => (NodeKind::FieldAccessor(*k), node.ty.clone()), NodeKind::GetField { rec, field } => { let rec_opt = self.visit_node(rec.clone(), sub, path); // Constant folding for Field Access if let NodeKind::Constant(Value::Record(layout, values)) = &rec_opt.kind && let Some(idx) = layout.index_of(*field) { return Rc::new(folder.make_constant_node(values[idx].clone(), node)); } if Rc::ptr_eq(&rec_opt, rec) { return node_rc; } ( NodeKind::GetField { rec: rec_opt, field: *field, }, node.ty.clone(), ) } NodeKind::Assign { target, value, info } => { let addr = info.addr; let value_opt = self.visit_node(value.clone(), sub, path); if let Some(addr) = addr { if let NodeKind::Constant(val) = &value_opt.kind { sub.add_value(addr, val.clone()); } else { sub.remove_value(&addr); } } // Remap target addresses without constant folding (targets are lvalues) let target_opt = Self::remap_pattern(target, sub); if Rc::ptr_eq(&value_opt, value) && Rc::ptr_eq(&target_opt, target) { return node_rc; } let new_info = if let Some(addr) = addr { AssignBinding { addr: Some(sub.map_address(addr)), } } else { info.clone() }; ( NodeKind::Assign { target: target_opt, value: value_opt, info: new_info, }, node.ty.clone(), ) } NodeKind::Def { pattern, value, info, } => { let value_opt = self.visit_node(value.clone(), sub, path); // Extract addr from the ORIGINAL pattern (before visiting) let addr = Self::extract_def_addr(pattern); if let Some(addr) = addr { if let Address::Local(slot) = addr && !info.captured_by.is_empty() { sub.captured_slots.insert(slot); } if let NodeKind::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.root_purity { let mut pr = purity_rc.borrow_mut(); let idx = global_index.0 as usize; if idx < pr.len() { pr[idx] = value_opt.ty.purity; } } } // Remap pattern addresses without constant folding (patterns are lvalues) let pattern_opt = Self::remap_pattern(pattern, sub); if Rc::ptr_eq(&value_opt, value) && Rc::ptr_eq(&pattern_opt, pattern) { return node_rc; } ( NodeKind::Def { pattern: pattern_opt, value: value_opt, info: info.clone(), }, node.ty.clone(), ) } NodeKind::Call { callee, args } => { let callee_opt = self.visit_node(callee.clone(), sub, path); let args_opt = self.visit_node(args.clone(), sub, path); if self.enabled { // Constant folding for Call { callee: FieldAccessor, args: Tuple[1] } // (field access on a constant record) if let NodeKind::FieldAccessor(k) = &callee_opt.kind && let NodeKind::Tuple { elements } = &args_opt.kind && elements.len() == 1 && let NodeKind::Constant(Value::Record(layout, values)) = &elements[0].kind && let Some(idx) = layout.index_of(*k) { return Rc::new(folder.make_constant_node(values[idx].clone(), node)); } let mut arg_nodes = Vec::new(); self.flatten_tuple(args_opt.clone(), &mut arg_nodes); if let NodeKind::Lambda { params, body, info: lambda_info, } = &callee_opt.kind && lambda_info.upvalues.is_empty() && lambda_info.positional_count.is_some() && path.inlining_depth < 5 && !callee_opt.ty.is_recursive && path.enter_lambda(&callee_opt.identity) { path.inlining_depth += 1; let collapsed = self.try_inline(params, &arg_nodes, body, sub, path, None); path.inlining_depth -= 1; path.exit_lambda(&callee_opt.identity); if let Some(res) = collapsed { return res; } } if let NodeKind::Identifier { binding: IdentifierBinding::Reference(Address::Global(idx)), .. } = &callee_opt.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 NodeKind::Lambda { params, body, info: lambda_info, } = &lambda_node.kind && lambda_info.upvalues.is_empty() && lambda_info.positional_count.is_some() && !lambda_node.ty.is_recursive { path.inlining_stack.insert(*idx); path.inlining_depth += 1; let collapsed = self.try_inline(params, &arg_nodes, body, sub, path, None); path.inlining_depth -= 1; path.inlining_stack.remove(idx); if let Some(res) = collapsed { return res; } } } if let NodeKind::Constant(Value::Object(ref obj)) = callee_opt.kind && path.inlining_depth < 5 && let Some(closure) = obj.as_any().downcast_ref::() && (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 = sub.new_for_inlining(); 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; if let NodeKind::Lambda { params, .. } = &closure.function_node.kind { let collapsed = self.try_inline( params, &arg_nodes, &closure.function_node, sub, path, Some(closure_sub), ); path.inlining_depth -= 1; path.exit_lambda(&closure.function_node.identity); if let Some(res) = collapsed { return res; } } else { path.inlining_depth -= 1; path.exit_lambda(&closure.function_node.identity); } } if let Some(folded) = folder.try_fold_pure(&callee_opt, &arg_nodes) { return Rc::new(folded); } } if Rc::ptr_eq(&callee_opt, callee) && Rc::ptr_eq(&args_opt, args) { return node_rc; } ( NodeKind::Call { callee: callee_opt, args: args_opt, }, node.ty.clone(), ) } NodeKind::Again { args } => { let args_opt = self.visit_node(args.clone(), sub, path); if Rc::ptr_eq(&args_opt, args) { return node_rc; } ( NodeKind::Again { args: args_opt, }, node.ty.clone(), ) } NodeKind::If { cond, then_br, else_br, } => { let cond_opt = self.visit_node(cond.clone(), sub, path); if self.enabled && let NodeKind::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 Rc::new(folder.make_nop_node(node)); } } let then_br_opt = self.visit_node(then_br.clone(), sub, path); let else_br_opt = else_br .as_ref() .map(|e| self.visit_node(e.clone(), sub, path)); let unchanged = Rc::ptr_eq(&cond_opt, cond) && Rc::ptr_eq(&then_br_opt, then_br) && match (&else_br_opt, else_br) { (Some(a), Some(b)) => Rc::ptr_eq(a, b), (None, None) => true, _ => false, }; if unchanged { return node_rc; } ( NodeKind::If { cond: cond_opt, then_br: then_br_opt, else_br: else_br_opt, }, node.ty.clone(), ) } NodeKind::Pipe { inputs, lambda, } => { let mut o_inputs = Vec::with_capacity(inputs.len()); let mut inputs_changed = false; for input in inputs { let opt = self.visit_node(input.clone(), sub, path); if !Rc::ptr_eq(&opt, input) { inputs_changed = true; } o_inputs.push(opt); } let o_lambda = self.visit_node(lambda.clone(), sub, path); if !inputs_changed && Rc::ptr_eq(&o_lambda, lambda) { return node_rc; } ( NodeKind::Pipe { inputs: o_inputs, lambda: o_lambda, }, node.ty.clone(), ) } NodeKind::Block { 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); let mut block_changed = false; for (i, e) in exprs.iter().enumerate() { let is_last = i == last_idx; if self.enabled && !is_last { let removable = match &e.kind { NodeKind::Def { pattern, value, info: def_info } => { let addr = Self::extract_def_addr(pattern); if let Some(addr) = addr { !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, NodeKind::Lambda { .. })) } else if def_info.captured_by.is_empty() { // Destructuring def: safe to remove only when no closure // captures any binding and every bound slot is unused. let mut addrs = Vec::new(); collect_pattern_addrs(pattern, &mut addrs); !addrs.is_empty() && addrs.iter().all(|a| { !info.is_used(a) && if let Address::Local(slot) = a { !sub.captured_slots.contains(slot) } else { true } }) && (value.ty.purity >= Purity::SideEffectFree || matches!(value.kind, NodeKind::Lambda { .. })) } else { false } } NodeKind::Assign { info: assign_info, value, .. } => { if let Some(addr) = assign_info.addr { !info.is_used(&addr) && (if let Address::Local(slot) = addr { !sub.captured_slots.contains(&slot) } else { true }) && value.ty.purity >= Purity::SideEffectFree } else { false } } _ => e.ty.purity >= Purity::SideEffectFree, }; if removable { block_changed = true; continue; } } let unmapped_addr = if let NodeKind::Def { pattern, .. } = &e.kind { Self::extract_def_addr(pattern) } else { None }; let opt = self.visit_node(e.clone(), sub, path); if let NodeKind::Def { value, .. } = &opt.kind && let Some(orig_addr) = unmapped_addr && !info.assigned.contains(&orig_addr) { let mut core_value = value.as_ref(); while let NodeKind::Expansion { expanded, .. } = &core_value.kind { core_value = expanded.as_ref(); } if let NodeKind::Constant(val) = &core_value.kind { sub.add_value(orig_addr, val.clone()); } else if let NodeKind::Lambda { info: lambda_info, .. } = &core_value.kind && lambda_info.upvalues.is_empty() { sub.add_ast_substitution(orig_addr, core_value.clone()); } } if self.enabled && matches!(opt.kind, NodeKind::Nop) && !is_last { block_changed = true; continue; } if !Rc::ptr_eq(&opt, e) { block_changed = true; } new_exprs.push(opt); } if !block_changed { return node_rc; } if self.enabled { if new_exprs.is_empty() { return Rc::new(folder.make_nop_node(node)); } else if new_exprs.len() == 1 { return new_exprs.pop().unwrap(); } } (NodeKind::Block { exprs: new_exprs }, node.ty.clone()) } NodeKind::Lambda { params, body, info: lambda_info, } => { let mut usage_info = UsageInfo::default(); usage_info.collect(node); let mut new_upvalues = Vec::new(); let mut mapping = Vec::new(); let mut next_inner_subs = sub.new_inner(); next_inner_subs.assigned = usage_info.assigned; let mut upvalues_changed = false; for (old_idx, capture_addr) in lambda_info.upvalues.iter().enumerate() { let mut inlined_val = None; let mut inlined_ast = None; if !sub.assigned.contains(capture_addr) { if let Some(val) = sub.get_value(capture_addr) { inlined_val = Some(val.clone()); } else if let Some(ast) = sub.ast_substitutions.get(capture_addr) { inlined_ast = Some(Rc::clone(ast)); } } 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); upvalues_changed = true; } else if let Some(ast) = inlined_ast { next_inner_subs .add_ast_substitution(Address::Upvalue(UpvalueIdx(old_idx as u32)), (*ast).clone()); mapping.push(None); upvalues_changed = true; } else { mapping.push(Some(new_upvalues.len() as u32)); let mapped = sub.map_address(*capture_addr); if mapped != *capture_addr { upvalues_changed = true; } new_upvalues.push(mapped); } } let params_opt = self.visit_node(params.clone(), &mut next_inner_subs, path); let body_opt = self.visit_node(body.clone(), &mut next_inner_subs, path); let reindexed_body = if new_upvalues.len() != lambda_info.upvalues.len() { sub.reindex_upvalues(body_opt, &mapping) } else { body_opt }; if !upvalues_changed && Rc::ptr_eq(¶ms_opt, params) && Rc::ptr_eq(&reindexed_body, body) { return node_rc; } ( NodeKind::Lambda { params: params_opt, body: reindexed_body, info: LambdaBinding { upvalues: new_upvalues, positional_count: lambda_info.positional_count, }, }, node.ty.clone(), ) } NodeKind::Tuple { elements } => { let mut new_elements = Vec::with_capacity(elements.len()); let mut changed = false; for e in elements { let opt = self.visit_node(e.clone(), sub, path); if !Rc::ptr_eq(&opt, e) { changed = true; } new_elements.push(opt); } if !changed { return node_rc; } (NodeKind::Tuple { elements: new_elements }, node.ty.clone()) } NodeKind::Record { fields, layout } => { let mut mapped_fields = Vec::with_capacity(fields.len()); let mut changed = false; for (k, v) in fields { let v_opt = self.visit_node(v.clone(), sub, path); if !Rc::ptr_eq(&v_opt, v) { changed = true; } mapped_fields.push((k.clone(), v_opt)); } if self.enabled { let values: Vec<_> = mapped_fields.iter().map(|(_, v)| v.clone()).collect(); if let Some(folded) = folder.try_fold_record(layout, &values, node) { return Rc::new(folded); } } if !changed { return node_rc; } ( NodeKind::Record { fields: mapped_fields, layout: layout.clone(), }, node.ty.clone(), ) } NodeKind::Expansion { original_call, expanded, } => { path.inlining_depth += 1; let expanded_opt = self.visit_node(expanded.clone(), sub, path); path.inlining_depth -= 1; if Rc::ptr_eq(&expanded_opt, expanded) { return node_rc; } ( NodeKind::Expansion { original_call: original_call.clone(), expanded: expanded_opt, }, node.ty.clone(), ) } k => (k.clone(), node.ty.clone()), }; Rc::new(Node { identity: node.identity.clone(), kind: new_kind, ty: metrics, }) } /// Extracts the address from a Def pattern node (when it's a simple Identifier with Declaration binding). fn extract_def_addr(pattern: &AnalyzedNode) -> Option> { if let NodeKind::Identifier { binding: IdentifierBinding::Declaration { addr, .. }, .. } = &pattern.kind { Some(*addr) } else { None } } /// Remaps addresses in a pattern/target node without constant folding. /// Used for Assign targets where identifiers are references to existing variables /// that should not be replaced by their values. fn remap_pattern(node_rc: &Rc, sub: &mut SubstitutionMap) -> Rc { let node = &**node_rc; let new_kind = match &node.kind { NodeKind::Identifier { symbol, binding } => { let addr = match binding { IdentifierBinding::Reference(addr) => *addr, IdentifierBinding::Declaration { addr, .. } => *addr, }; let new_binding = match binding { IdentifierBinding::Reference(_) => { IdentifierBinding::Reference(sub.map_address(addr)) } IdentifierBinding::Declaration { kind, .. } => { IdentifierBinding::Declaration { addr: sub.map_address(addr), kind: *kind, } } }; NodeKind::Identifier { symbol: symbol.clone(), binding: new_binding, } } NodeKind::Tuple { elements } => { let new_elements = elements .iter() .map(|e| Self::remap_pattern(e, sub)) .collect(); NodeKind::Tuple { elements: new_elements, } } _ => return node_rc.clone(), }; Rc::new(Node { identity: node.identity.clone(), kind: new_kind, ty: node.ty.clone(), }) } fn flatten_tuple(&self, node: Rc, into: &mut Vec>) { match &node.kind { NodeKind::Tuple { elements } => { for el in elements { self.flatten_tuple(el.clone(), into); } } NodeKind::Nop => {} _ => into.push(node), } } }