use crate::ast::compiler::bound_nodes::{Address, BoundKind, TypedNode}; use crate::ast::nodes::Node; use crate::ast::types::{StaticType, Value}; use crate::ast::vm::Closure; use std::cell::RefCell; use std::collections::{HashMap, HashSet}; use std::rc::Rc; #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub enum Purity { Impure, // Has side effects (e.g. Set, Print) SideEffectFree, // No side effects, but not deterministic (e.g. now(), random()) Pure, // No side effects and deterministic (e.g. sin(x), 1 + 2) } /// The Optimizer performs Phase 2 (Cracking), Phase 2.5/2.6 (Aggressive Collapsing & DCE) /// and Phase 4 (Global Inlining). pub struct Optimizer { pub enabled: bool, max_passes: usize, pub globals: Option>>>, pub global_purity: Option>>>, pub lambda_registry: Option>>>, } /// Global state for tracking the optimization path (recursion, depth). struct PathTracker { inlining_depth: usize, /// Stack of global indices to detect recursion in named functions. inlining_stack: HashSet, /// Stack of node identities to detect recursion in anonymous lambdas. identity_stack: HashSet, } 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>>) -> Self { self.globals = Some(globals); self } pub fn with_purity(mut self, purity: Rc>>) -> Self { self.global_purity = Some(purity); self } pub fn with_registry(mut self, registry: Rc>>) -> Self { self.lambda_registry = Some(registry); self } pub fn optimize(&self, node: TypedNode) -> TypedNode { 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: TypedNode, sub: &mut SubstitutionMap, path: &mut PathTracker, ) -> TypedNode { let (new_kind, new_ty) = match node.kind { BoundKind::Get { addr, name } => { let new_addr = match addr { Address::Local(slot) => { if !sub.assigned_locals.contains(&slot) { if let Some(inlined_node) = sub.ast_locals.get(&slot) { return inlined_node.clone(); } if let Some(val) = sub.locals.get(&slot) { return Node { identity: node.identity, ty: val.static_type(), kind: BoundKind::Constant(val.clone()), }; } } sub.used_slots.insert(slot); Address::Local(sub.map_slot(slot)) } Address::Upvalue(idx) => { if let Some(val) = sub.upvalues.get(&idx) { return Node { identity: node.identity, ty: val.static_type(), kind: BoundKind::Constant(val.clone()), }; } Address::Upvalue(idx) } Address::Global(idx) => { if !sub.assigned_globals.contains(&idx) { if let Some(val) = sub.globals.get(&idx) { return Node { identity: node.identity, ty: val.static_type(), kind: BoundKind::Constant(val.clone()), }; } if let Some(globals_rc) = &self.globals { let globals = globals_rc.borrow(); if let Some(val) = globals.get(idx as usize) && self.is_inlinable_value(val) { return Node { identity: node.identity, ty: val.static_type(), kind: BoundKind::Constant(val.clone()), }; } } } sub.used_globals.insert(idx); Address::Global(idx) } }; ( BoundKind::Get { addr: new_addr, name, }, node.ty, ) } BoundKind::Parameter { name, slot } => { let new_slot = sub.map_slot(slot); ( BoundKind::Parameter { name, slot: new_slot, }, node.ty, ) } BoundKind::Set { addr, value } => { let value = Box::new(self.visit_node(*value, sub, path)); let new_addr = match addr { Address::Local(slot) => { if let BoundKind::Constant(val) = &value.kind { sub.add_local(slot, val.clone()); } else { sub.locals.remove(&slot); } Address::Local(sub.map_slot(slot)) } Address::Global(idx) => { if let BoundKind::Constant(val) = &value.kind { sub.add_global(idx, val.clone()); } else { sub.globals.remove(&idx); } Address::Global(idx) } _ => addr, }; ( BoundKind::Set { addr: new_addr, value, }, node.ty, ) } BoundKind::Call { callee, args } => { let callee = self.visit_node(*callee, sub, path); let args = self.visit_node(*args, sub, path); if self.enabled { // Case 1: Beta-Reduction for Lambda Literals if let BoundKind::Lambda { params, body, upvalues, positional_count, } = &callee.kind && upvalues.is_empty() && positional_count.is_some() && path.inlining_depth < 5 { if 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; } } } // Case 1.5: Beta-Reduction for Global Registry functions (from Registry) 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() { // RECURSION CHECK: Don't inline if the function calls itself let mut used_in_body = HashSet::new(); let mut assigned_in_body = HashSet::new(); let mut used_globals_in_body = HashSet::new(); let mut assigned_globals_in_body = HashSet::new(); self.collect_usage( body, &mut used_in_body, &mut used_globals_in_body, &mut assigned_in_body, &mut assigned_globals_in_body, ); if !used_globals_in_body.contains(idx) { 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; } } } } // Case 2: Cracking and Inlining for Constant Closures if let BoundKind::Constant(Value::Object(ref obj)) = callee.kind && path.inlining_depth < 5 && let Some(closure) = obj.as_any().downcast_ref::() && (closure.upvalues.is_empty() || self.purity_of(&closure.function_node) >= Purity::SideEffectFree) { let mut closure_sub = SubstitutionMap::new(); for (i, cell) in closure.upvalues.iter().enumerate() { closure_sub.add_upvalue(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; if let Some(res) = collapsed { return res; } } if let Some(folded) = self.try_fold_pure(&callee, &args) { return folded; } } if self.enabled && let BoundKind::Constant(Value::Object(ref obj)) = callee.kind && path.inlining_depth < 5 && let Some(closure) = obj.as_any().downcast_ref::() && (closure.upvalues.is_empty() || self.purity_of(&closure.function_node) >= Purity::SideEffectFree) { let mut closure_sub = SubstitutionMap::new(); for (i, cell) in closure.upvalues.iter().enumerate() { closure_sub.add_upvalue(i as u32, cell.borrow().clone()); } path.inlining_depth += 1; let inlined_body = self.visit_node((*closure.function_node).clone(), &mut closure_sub, path); path.inlining_depth -= 1; 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, }; } ( 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, sub, path); if self.enabled && let BoundKind::Constant(ref val) = cond.kind { if val.is_truthy() { return self.visit_node(*then_br, sub, path); } else if let Some(else_node) = else_br { return self.visit_node(*else_node, sub, path); } else { return Node { identity: node.identity, kind: BoundKind::Nop, ty: StaticType::Void, }; } } let then_br = Box::new(self.visit_node(*then_br, sub, path)); let else_br = else_br.map(|e| Box::new(self.visit_node(*e, sub, path))); ( BoundKind::If { cond: Box::new(cond), then_br, else_br, }, node.ty, ) } BoundKind::Block { exprs } => { let mut used_in_block = HashSet::new(); let mut used_globals_in_block = HashSet::new(); let mut assigned_locals_in_block = HashSet::new(); let mut assigned_globals_in_block = HashSet::new(); if !exprs.is_empty() { for e in &exprs { self.collect_usage( e, &mut used_in_block, &mut used_globals_in_block, &mut assigned_locals_in_block, &mut assigned_globals_in_block, ); } if let Some(last) = exprs.last() { match &last.kind { BoundKind::Get { addr: Address::Local(slot), .. } => { used_in_block.insert(*slot); } BoundKind::Get { addr: Address::Global(idx), .. } => { used_globals_in_block.insert(*idx); } _ => {} } } } // IMPORTANT: Propagate block assignments to the outer map so we don't // incorrectly inline these variables in the rest of the script. for slot in &assigned_locals_in_block { sub.assigned_locals.insert(*slot); } for idx in &assigned_globals_in_block { sub.assigned_globals.insert(*idx); } let mut new_exprs = Vec::with_capacity(exprs.len()); let last_idx = exprs.len().saturating_sub(1); for (i, e) in exprs.into_iter().enumerate() { let is_last = i == last_idx; if self.enabled && !is_last { let removable = match &e.kind { BoundKind::DefLocal { slot, value, .. } => { !used_in_block.contains(slot) && !sub.captured_slots.contains(slot) && self.purity_of(value) >= Purity::SideEffectFree } BoundKind::DefGlobal { global_index, value, .. } => { !used_globals_in_block.contains(global_index) && (self.purity_of(value) >= Purity::SideEffectFree || matches!(value.kind, BoundKind::Lambda { .. })) } BoundKind::Set { addr: Address::Local(slot), value, .. } => { !used_in_block.contains(slot) && !sub.captured_slots.contains(slot) && self.purity_of(value) >= Purity::SideEffectFree } _ => self.purity_of(&e) >= Purity::SideEffectFree, }; if removable { continue; } } let opt = self.visit_node(e, 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 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) } 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 used_in_lambda = HashSet::new(); let mut used_globals_in_lambda = HashSet::new(); let mut assigned_locals_in_lambda = HashSet::new(); let mut assigned_globals_in_lambda = HashSet::new(); self.collect_usage( &node, &mut used_in_lambda, &mut used_globals_in_lambda, &mut assigned_locals_in_lambda, &mut assigned_globals_in_lambda, ); // Track captures from outer scope that are assigned in this lambda for idx in &assigned_globals_in_lambda { sub.assigned_globals.insert(*idx); } let mut new_upvalues = Vec::new(); let mut mapping = Vec::new(); let mut next_inner_subs = SubstitutionMap::new(); next_inner_subs.assigned_locals = assigned_locals_in_lambda; next_inner_subs.assigned_globals = assigned_globals_in_lambda; for (old_idx, capture_addr) in original_upvalues.iter().enumerate() { let mut inlined_val = None; match capture_addr { Address::Local(slot) => { if let Some(val) = sub.locals.get(slot) { inlined_val = Some(val.clone()); } sub.captured_slots.insert(*slot); } Address::Upvalue(idx) => { if let Some(val) = sub.upvalues.get(idx) { inlined_val = Some(val.clone()); } } _ => {} } if let Some(val) = inlined_val { next_inner_subs.add_upvalue(old_idx as u32, val); mapping.push(None); } else { mapping.push(Some(new_upvalues.len() as u32)); let new_addr = if let Address::Local(parent_slot) = capture_addr { Address::Local(sub.map_slot(*parent_slot)) } else { *capture_addr }; new_upvalues.push(new_addr); } } // 1. Visit parameters to determine their new slots let params_node = self.visit_node(params.as_ref().clone(), &mut next_inner_subs, path); // 2. IMPORTANT: Pre-register parameter slots in the inner substitution map // so they are mapped 1:1 and don't get reassigned to higher slots in the body. self.collect_parameter_slots(¶ms_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, ) } BoundKind::DefLocal { name, slot, value, captured_by, } => { if !captured_by.is_empty() { sub.captured_slots.insert(slot); } let value = Box::new(self.visit_node(*value, sub, path)); if let BoundKind::Constant(val) = &value.kind { sub.add_local(slot, val.clone()); } else { sub.locals.remove(&slot); } if self.purity_of(&value) == Purity::Pure { sub.pure_slots.insert(slot); } else { sub.pure_slots.remove(&slot); } let new_slot = sub.map_slot(slot); ( BoundKind::DefLocal { name, slot: new_slot, value, captured_by, }, node.ty, ) } BoundKind::DefGlobal { name, global_index, value, } => { let value = Box::new(self.visit_node(*value, sub, path)); if let BoundKind::Constant(val) = &value.kind { sub.add_global(global_index, val.clone()); } let p = self.purity_of(&value); if p > Purity::Impure && let Some(purity_rc) = &self.global_purity { purity_rc.borrow_mut().insert(global_index, p); } ( BoundKind::DefGlobal { name, global_index, value, }, node.ty, ) } BoundKind::Tuple { elements } => { let elements = elements .into_iter() .map(|e| self.visit_node(e, sub, path)) .collect(); (BoundKind::Tuple { elements }, node.ty) } 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) } BoundKind::Expansion { original_call, bound_expanded, } => { path.inlining_depth += 1; let bound_expanded = Box::new(self.visit_node(*bound_expanded, sub, path)); path.inlining_depth -= 1; ( BoundKind::Expansion { original_call, bound_expanded, }, node.ty, ) } k => (k, node.ty), }; Node { identity: node.identity, kind: new_kind, ty: new_ty, } } fn collect_parameter_slots(&self, node: &TypedNode, sub: &mut SubstitutionMap) { match &node.kind { BoundKind::Parameter { slot, .. } => { // Identity mapping for parameters to keep them in their reserved frame slots sub.slot_mapping.insert(*slot, *slot); sub.next_slot = sub.next_slot.max(*slot + 1); } BoundKind::Tuple { elements } => { for el in elements { self.collect_parameter_slots(el, sub); } } _ => {} } } fn is_inlinable_value(&self, val: &Value) -> bool { 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.upvalues.is_empty() } else { false } } _ => false, } } fn purity_of(&self, node: &TypedNode) -> Purity { match &node.kind { BoundKind::Constant(_) | BoundKind::Parameter { .. } | BoundKind::Nop => Purity::Pure, // Defining a lambda is pure; only calling it may have side effects. BoundKind::Lambda { .. } => Purity::Pure, BoundKind::Get { addr, .. } => match addr { Address::Global(idx) => { if let Some(purity_rc) = &self.global_purity { *purity_rc.borrow().get(idx).unwrap_or(&Purity::Impure) } else { Purity::Impure } } _ => Purity::Pure, }, BoundKind::Tuple { elements } => elements .iter() .map(|e| self.purity_of(e)) .min() .unwrap_or(Purity::Pure), BoundKind::Record { fields } => fields .iter() .map(|(k, v)| self.purity_of(k).min(self.purity_of(v))) .min() .unwrap_or(Purity::Pure), BoundKind::If { cond, then_br, else_br, } => { let p = self.purity_of(cond).min(self.purity_of(then_br)); if let Some(e) = else_br { p.min(self.purity_of(e)) } else { p } } BoundKind::Block { exprs } => exprs .iter() .map(|e| self.purity_of(e)) .min() .unwrap_or(Purity::Pure), BoundKind::Expansion { bound_expanded, .. } => self.purity_of(bound_expanded), BoundKind::DefLocal { value, .. } => self.purity_of(value), BoundKind::Set { .. } => Purity::Impure, BoundKind::Call { callee, args } => { let callee_purity = match &callee.kind { BoundKind::Lambda { body, .. } => self.purity_of(body), BoundKind::Get { addr: Address::Global(idx), .. } => { if let Some(purity_rc) = &self.global_purity { *purity_rc.borrow().get(idx).unwrap_or(&Purity::Impure) } else { Purity::Impure } } BoundKind::Constant(Value::Function(_)) => Purity::Pure, BoundKind::Constant(Value::Object(obj)) => { if let Some(closure) = obj.as_any().downcast_ref::() { self.purity_of(&closure.function_node) } else { Purity::Impure } } _ => Purity::Impure, }; callee_purity.min(self.purity_of(args)) } _ => Purity::Impure, } } fn try_fold_pure(&self, callee: &TypedNode, args: &TypedNode) -> Option { let temp_call = Node { identity: callee.identity.clone(), ty: StaticType::Any, kind: BoundKind::Call { callee: Box::new(callee.clone()), args: Box::new(args.clone()), }, }; // Folding requires the function AND arguments to be truly Pure (deterministic). if self.purity_of(&temp_call) < Purity::Pure { return None; } let mut arg_nodes = Vec::new(); self.flatten_typed_tuple(args, &mut arg_nodes); let mut arg_values = Vec::with_capacity(arg_nodes.len()); for node in arg_nodes { if let BoundKind::Constant(val) = node.kind { arg_values.push(val); } else { return None; } } let func_val = match &callee.kind { BoundKind::Get { addr: Address::Global(idx), .. } => self.globals.as_ref()?.borrow().get(*idx as usize)?.clone(), BoundKind::Constant(val) => val.clone(), _ => return None, }; let result = match func_val { Value::Function(f) => f(arg_values), _ => return None, }; Some(Node { identity: callee.identity.clone(), ty: result.static_type(), kind: BoundKind::Constant(result), }) } fn try_beta_reduce_with_sub( &self, params: &TypedNode, args: &TypedNode, body: TypedNode, sub: &mut SubstitutionMap, path: &mut PathTracker, ) -> Option { let mut arg_vals = Vec::new(); self.flatten_typed_tuple(args, &mut arg_vals); let mut slot_index = 0; self.map_params_to_args(params, &arg_vals, &mut slot_index, sub); // Safety: Only proceed if we mapped EXACTLY the number of arguments provided. // And if we actually have parameters to bind (or if it's a 0-arg call). if slot_index != arg_vals.len() { return None; } if sub.locals.is_empty() && sub.ast_locals.is_empty() && sub.upvalues.is_empty() && !arg_vals.is_empty() { return None; } Some(self.visit_node(body, sub, path)) } fn flatten_typed_tuple(&self, node: &TypedNode, into: &mut Vec) { if let BoundKind::Tuple { elements } = &node.kind { for el in elements { self.flatten_typed_tuple(el, into); } } else if !matches!(node.kind, BoundKind::Nop) { into.push(node.clone()); } } fn map_params_to_args( &self, pattern: &TypedNode, args: &[TypedNode], offset: &mut usize, sub: &mut SubstitutionMap, ) { match &pattern.kind { BoundKind::Parameter { slot, .. } => { if let Some(arg) = args.get(*offset) { if let BoundKind::Constant(val) = &arg.kind { sub.add_local(*slot, val.clone()); } else if let BoundKind::Lambda { upvalues, .. } = &arg.kind { if upvalues.is_empty() { sub.add_ast_local(*slot, arg.clone()); } } else if let BoundKind::Get { addr: Address::Global(_), .. } = &arg.kind { sub.add_ast_local(*slot, arg.clone()); } } // Parameters in beta-reduction are also identity mapped if not inlined sub.slot_mapping.insert(*slot, *slot); sub.next_slot = sub.next_slot.max(*slot + 1); *offset += 1; } BoundKind::Tuple { elements } => { for el in elements { self.map_params_to_args(el, args, offset, sub); } } _ => {} } } fn collect_usage( &self, node: &TypedNode, used_locals: &mut HashSet, used_globals: &mut HashSet, assigned_locals: &mut HashSet, assigned_globals: &mut HashSet, ) { match &node.kind { BoundKind::Constant(v) => { if let Value::Object(obj) = v && let Some(closure) = obj.as_any().downcast_ref::() { self.collect_usage( &closure.function_node, used_locals, used_globals, assigned_locals, assigned_globals, ); } } BoundKind::Get { addr, .. } => match addr { Address::Local(slot) => { used_locals.insert(*slot); } Address::Global(idx) => { used_globals.insert(*idx); } _ => {} }, BoundKind::Set { addr, value } => { match addr { Address::Local(slot) => { assigned_locals.insert(*slot); } Address::Global(idx) => { assigned_globals.insert(*idx); } _ => {} } self.collect_usage( value, used_locals, used_globals, assigned_locals, assigned_globals, ); } BoundKind::Lambda { params, body, upvalues, .. } => { self.collect_usage( params, used_locals, used_globals, assigned_locals, assigned_globals, ); self.collect_usage( body, used_locals, used_globals, assigned_locals, assigned_globals, ); for addr in upvalues { match addr { Address::Local(slot) => { used_locals.insert(*slot); } Address::Global(idx) => { used_globals.insert(*idx); } _ => {} } } } BoundKind::Block { exprs } => { for e in exprs { self.collect_usage( e, used_locals, used_globals, assigned_locals, assigned_globals, ); } } BoundKind::If { cond, then_br, else_br, } => { self.collect_usage( cond, used_locals, used_globals, assigned_locals, assigned_globals, ); self.collect_usage( then_br, used_locals, used_globals, assigned_locals, assigned_globals, ); if let Some(e) = else_br { self.collect_usage( e, used_locals, used_globals, assigned_locals, assigned_globals, ); } } BoundKind::Call { callee, args } => { self.collect_usage( callee, used_locals, used_globals, assigned_locals, assigned_globals, ); self.collect_usage( args, used_locals, used_globals, assigned_locals, assigned_globals, ); } BoundKind::Tuple { elements } => { for e in elements { self.collect_usage( e, used_locals, used_globals, assigned_locals, assigned_globals, ); } } BoundKind::Record { fields } => { for (k, v) in fields { self.collect_usage( k, used_locals, used_globals, assigned_locals, assigned_globals, ); self.collect_usage( v, used_locals, used_globals, assigned_locals, assigned_globals, ); } } BoundKind::DefLocal { value, .. } | BoundKind::DefGlobal { value, .. } => { self.collect_usage( value, used_locals, used_globals, assigned_locals, assigned_globals, ); } BoundKind::Expansion { bound_expanded, .. } => { self.collect_usage( bound_expanded, used_locals, used_globals, assigned_locals, assigned_globals, ); } _ => {} } } } /// Helper for transitively inlining values and cleaning up capture lists. struct SubstitutionMap { locals: HashMap, upvalues: HashMap, globals: HashMap, ast_locals: HashMap, slot_mapping: HashMap, assigned_locals: HashSet, assigned_globals: HashSet, next_slot: u32, used_slots: HashSet, pure_slots: HashSet, used_globals: HashSet, captured_slots: HashSet, } impl SubstitutionMap { fn new() -> Self { Self { locals: HashMap::new(), upvalues: HashMap::new(), globals: HashMap::new(), ast_locals: HashMap::new(), slot_mapping: HashMap::new(), assigned_locals: HashSet::new(), assigned_globals: HashSet::new(), next_slot: 0, used_slots: HashSet::new(), pure_slots: HashSet::new(), used_globals: HashSet::new(), captured_slots: HashSet::new(), } } fn add_ast_local(&mut self, slot: u32, node: TypedNode) { self.ast_locals.insert(slot, node); } fn map_slot(&mut self, old_slot: u32) -> u32 { if let Some(&new_slot) = self.slot_mapping.get(&old_slot) { return new_slot; } let new_slot = self.next_slot; self.slot_mapping.insert(old_slot, new_slot); self.next_slot += 1; new_slot } fn add_local(&mut self, slot: u32, val: Value) { self.locals.insert(slot, val); } fn add_upvalue(&mut self, idx: u32, val: Value) { self.upvalues.insert(idx, val); } fn add_global(&mut self, idx: u32, val: Value) { self.globals.insert(idx, val); } fn reindex_upvalues(&self, node: TypedNode, mapping: &[Option]) -> TypedNode { let (new_kind, new_ty) = match node.kind { BoundKind::Get { addr: Address::Upvalue(idx), name, } => { if let Some(res) = mapping.get(idx as usize) { match res { Some(new_idx) => ( BoundKind::Get { addr: Address::Upvalue(*new_idx), name, }, node.ty, ), None => ( BoundKind::Get { addr: Address::Upvalue(idx), name, }, node.ty, ), } } else { ( BoundKind::Get { addr: Address::Upvalue(idx), name, }, node.ty, ) } } BoundKind::Lambda { params, upvalues, body, positional_count, } => { let mut next_upvalues = Vec::new(); for addr in upvalues { if let Address::Upvalue(idx) = addr && let Some(res) = mapping.get(idx as usize) { if let Some(new_idx) = res { next_upvalues.push(Address::Upvalue(*new_idx)); } continue; } next_upvalues.push(addr); } ( BoundKind::Lambda { params, upvalues: next_upvalues, body, positional_count, }, node.ty, ) } BoundKind::If { cond, then_br, else_br, } => { let cond = Box::new(self.reindex_upvalues(*cond, mapping)); let then_br = Box::new(self.reindex_upvalues(*then_br, mapping)); let else_br = else_br.map(|e| Box::new(self.reindex_upvalues(*e, mapping))); ( BoundKind::If { cond, then_br, else_br, }, node.ty, ) } BoundKind::Block { exprs } => { let exprs = exprs .into_iter() .map(|e| self.reindex_upvalues(e, mapping)) .collect(); (BoundKind::Block { exprs }, node.ty) } BoundKind::Call { callee, args } => { let callee = Box::new(self.reindex_upvalues(*callee, mapping)); let args = Box::new(self.reindex_upvalues(*args, mapping)); (BoundKind::Call { callee, args }, node.ty) } BoundKind::DefLocal { name, slot, value, captured_by, } => { let value = Box::new(self.reindex_upvalues(*value, mapping)); ( BoundKind::DefLocal { name, slot, value, captured_by, }, node.ty, ) } BoundKind::DefGlobal { name, global_index, value, } => { let value = Box::new(self.reindex_upvalues(*value, mapping)); ( BoundKind::DefGlobal { name, global_index, value, }, node.ty, ) } BoundKind::Set { addr, value } => { let value = Box::new(self.reindex_upvalues(*value, mapping)); (BoundKind::Set { addr, value }, node.ty) } BoundKind::Tuple { elements } => { let elements = elements .into_iter() .map(|e| self.reindex_upvalues(e, mapping)) .collect(); (BoundKind::Tuple { elements }, node.ty) } BoundKind::Record { fields } => { let fields = fields .into_iter() .map(|(k, v)| { ( self.reindex_upvalues(k, mapping), self.reindex_upvalues(v, mapping), ) }) .collect(); (BoundKind::Record { fields }, node.ty) } k => (k, node.ty), }; Node { identity: node.identity, kind: new_kind, ty: new_ty, } } } #[cfg(test)] mod tests { use crate::ast::environment::Environment; fn get_optimized_dump(source: &str, enabled: bool) -> String { let mut env = Environment::new(); env.optimization = enabled; env.dump_ast(source) .expect("Compilation failed during test") } #[test] fn test_opt_folding() { let dump = get_optimized_dump("(+ 10 20)", true); insta::assert_snapshot!(dump); } #[test] fn test_opt_folding_complex() { let dump_float = get_optimized_dump("(+ 1.5 2.5)", true); assert!(dump_float.contains("Constant: 4")); let dump_text = get_optimized_dump("(+ \"hello \" \"world\")", true); assert!(dump_text.contains("Constant: \"hello world\"")); let dump_date = get_optimized_dump("(date \"2023-01-01\")", true); assert!(dump_date.contains("Constant: #2023-01-01 00:00:00#")); let dump_cmp = get_optimized_dump("(> 10 5)", true); assert!(dump_cmp.contains("Constant: true")); } #[test] fn test_opt_local_inlining() { let source = "(fn [] (do (def x 10) (+ x 5)))"; let dump = get_optimized_dump(source, true); insta::assert_snapshot!(dump); } #[test] fn test_opt_dce_unused() { let source = "(fn [] (do (def x 10) 42))"; let dump = get_optimized_dump(source, true); insta::assert_snapshot!(dump); } #[test] fn test_opt_assignment_safety() { let source = "(fn [] (do (def x 10) (assign x 20) x))"; let dump = get_optimized_dump(source, true); insta::assert_snapshot!(dump); } #[test] fn test_opt_lambda_cracking() { let source = "(fn [] (do (def x 10) (fn [] x)))"; let dump = get_optimized_dump(source, true); insta::assert_snapshot!(dump); } #[test] fn test_opt_global_unused_dce() { let source = "(do (def f (fn [x] x)) 42)"; let dump = get_optimized_dump(source, true); assert!( !dump.contains("DefGlobal"), "Unused global helper 'f' should be optimized away. Dump: \n{}", dump ); assert!(dump.contains("Constant: 42")); } #[test] fn test_opt_local_recursion_safety() { // (fn [] (do (def f (fn [x] (f x))) (f 1))) // Hier sollte der Optimizer 'f' nicht unendlich oft in sich selbst inlinen. let source = "(fn [] (do (def f (fn [x] (f x))) (f 1)))"; let dump = get_optimized_dump(source, true); // Der Dump sollte 'Call' auf 'f' (oder Slot-Get) noch enthalten, // statt 5 Ebenen tief entfaltet zu sein. assert!(dump.contains("Call"), "Recursive call should remain as a call. Dump: \n{}", dump); assert!(!dump.contains("- Capturer:"), "Should not be over-optimized. Dump: \n{}", dump); } }