From e87f7232e6f78e81630d478dcc5fe50bd14a9797 Mon Sep 17 00:00:00 2001 From: Michael Schimmel Date: Sun, 22 Feb 2026 00:41:02 +0100 Subject: [PATCH] Refactor: Improve optimizer and DCE This commit introduces several improvements to the AST optimizer and dead code elimination (DCE) logic: - **Address Mapping:** The `Get` and `Set` operations now correctly map local slots and global indices using the substitution map, ensuring that remapped variables are handled properly. - **Parameter Slot Mapping:** Parameters are now explicitly mapped to their correct slots within the `Lambda` node, preventing potential slot reassignments in the function body. - **Pure Function Analysis:** The `is_pure` function has been refined to accurately identify pure expressions, including calls to pure functions and stateless closures. - **Dead Code Elimination:** DCE logic in `BoundKind::Block` has been enhanced to remove unused local definitions and pure, non-essential expressions more effectively. Global DCE is also improved. - **Constant Folding:** `try_fold_pure` is now more robust in folding pure function calls with constant arguments, reducing redundant computations. - **Beta Reduction:** `try_beta_reduce` has been updated to avoid reducing if the body contains `DefLocal`, ensuring correctness. - **Global Purity Tracking:** A `global_purity` map is introduced to track the purity of global values, enabling better optimization of global accesses. - **Inlinable Value Check:** The `is_inlinable_value` check now correctly identifies stateless closures for inlining. - **Optimization Level:** The default `optimization_level` is set to 2 to enable more aggressive optimizations. --- src/ast/compiler/optimizer.rs | 328 ++++++++++++++++++---------------- src/ast/environment.rs | 2 +- 2 files changed, 175 insertions(+), 155 deletions(-) diff --git a/src/ast/compiler/optimizer.rs b/src/ast/compiler/optimizer.rs index 9cf7d4e..080b2ec 100644 --- a/src/ast/compiler/optimizer.rs +++ b/src/ast/compiler/optimizer.rs @@ -50,7 +50,7 @@ impl Optimizer { fn visit_node(&self, node: TypedNode, sub: &mut SubstitutionMap) -> TypedNode { let (new_kind, new_ty) = match node.kind { BoundKind::Get { addr, name } => { - match addr { + let new_addr = match addr { Address::Local(slot) => { if let Some(val) = sub.locals.get(&slot) { return Node { @@ -60,6 +60,7 @@ impl Optimizer { }; } sub.used_slots.insert(slot); + Address::Local(sub.map_slot(slot)) } Address::Upvalue(idx) => { if let Some(val) = sub.upvalues.get(&idx) { @@ -69,10 +70,9 @@ impl Optimizer { kind: BoundKind::Constant(val.clone()), }; } + Address::Upvalue(idx) } Address::Global(idx) => { - // Phase 4: Global Inlining - // 1. Check script-local global substitution map if let Some(val) = sub.globals.get(&idx) { return Node { identity: node.identity, @@ -80,11 +80,9 @@ impl Optimizer { kind: BoundKind::Constant(val.clone()), }; } - // 2. Check environment global values if let Some(globals_rc) = &self.globals { let globals = globals_rc.borrow(); if let Some(val) = globals.get(idx as usize) { - // Only inline "stable" values: scalars and stateless functions if self.is_inlinable_value(val) { return Node { identity: node.identity, @@ -95,11 +93,41 @@ impl Optimizer { } } sub.used_globals.insert(idx); + Address::Global(idx) } - } - (BoundKind::Get { addr, name }, node.ty) + }; + (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)); + 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); let args = self.visit_node(*args, sub); @@ -178,44 +206,48 @@ impl Optimizer { }, BoundKind::Block { exprs } => { + let mut used_in_block = HashSet::new(); + if !exprs.is_empty() { + for e in &exprs { + self.collect_local_usage(e, &mut used_in_block); + } + if let Some(last) = exprs.last() { + if let BoundKind::Get { addr: Address::Local(slot), .. } = &last.kind { + used_in_block.insert(*slot); + } + } + } + let mut new_exprs = Vec::with_capacity(exprs.len()); - for e in exprs { + let last_idx = exprs.len().saturating_sub(1); + + for (i, e) in exprs.into_iter().enumerate() { + let is_last = i == last_idx; + + if self.level >= 2 && !is_last { + let removable = match &e.kind { + BoundKind::DefLocal { slot, value, .. } => { + !used_in_block.contains(slot) && !sub.captured_slots.contains(slot) && self.is_pure(value) + } + BoundKind::Set { addr: Address::Local(slot), value, .. } => { + !used_in_block.contains(slot) && !sub.captured_slots.contains(slot) && self.is_pure(value) + } + _ => self.is_pure(&e), + }; + + if removable { + continue; + } + } + let opt = self.visit_node(e, sub); - if self.level >= 2 && matches!(opt.kind, BoundKind::Nop) { + if self.level >= 2 && matches!(opt.kind, BoundKind::Nop) && !is_last { continue; } new_exprs.push(opt); } if self.level >= 2 { - if !new_exprs.is_empty() { - let last_idx = new_exprs.len() - 1; - let mut filtered = Vec::with_capacity(new_exprs.len()); - for (i, e) in new_exprs.into_iter().enumerate() { - if i == last_idx { - filtered.push(e); - continue; - } - - let removable = match &e.kind { - BoundKind::DefLocal { slot, .. } | BoundKind::Set { addr: Address::Local(slot), .. } => { - !sub.used_slots.contains(slot) && !sub.captured_slots.contains(slot) && self.is_pure(&e) - } - BoundKind::DefGlobal { global_index, .. } => { - // Phase 4: DCE for Globals - // A global can be removed if it was only used for inlining in this script. - !sub.used_globals.contains(global_index) && self.is_pure(&e) - } - _ => self.is_pure(&e), - }; - - if !removable { - filtered.push(e); - } - } - new_exprs = filtered; - } - if new_exprs.is_empty() { return Node { identity: node.identity, kind: BoundKind::Nop, ty: StaticType::Void }; } else if new_exprs.len() == 1 { @@ -250,13 +282,23 @@ impl Optimizer { mapping.push(None); } else { mapping.push(Some(new_upvalues.len() as u32)); - new_upvalues.push(*capture_addr); + 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); } } - let params = Rc::new(self.visit_node(params.as_ref().clone(), &mut SubstitutionMap::new())); - let body_node = self.visit_node((*body).clone(), &mut next_inner_subs); + // 1. Visit parameters to determine their new slots + let params_node = self.visit_node(params.as_ref().clone(), &mut SubstitutionMap::new()); + // 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).clone(), &mut next_inner_subs); let reindexed_body = if new_upvalues.len() != original_upvalues.len() { sub.reindex_upvalues(body_node, &mapping) } else { @@ -264,7 +306,7 @@ impl Optimizer { }; (BoundKind::Lambda { - params, + params: Rc::new(params_node), upvalues: new_upvalues, body: Rc::new(reindexed_body), positional_count @@ -281,45 +323,30 @@ impl Optimizer { } else { sub.locals.remove(&slot); } - (BoundKind::DefLocal { name, slot, value, captured_by }, node.ty) + + if self.is_pure(&value) { + 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)); - // Phase 4: Track script-local global definitions for inlining if let BoundKind::Constant(val) = &value.kind { sub.add_global(global_index, val.clone()); } - - // Purity Inference: Globals are pure if their value is pure if self.is_pure(&value) { if let Some(purity_rc) = &self.global_purity { purity_rc.borrow_mut().insert(global_index, true); } } - (BoundKind::DefGlobal { name, global_index, value }, node.ty) }, - BoundKind::Set { addr, value } => { - let value = Box::new(self.visit_node(*value, sub)); - match addr { - Address::Local(slot) => { - if let BoundKind::Constant(val) = &value.kind { - sub.add_local(slot, val.clone()); - } else { - sub.locals.remove(&slot); - } - } - Address::Global(idx) => { - if let BoundKind::Constant(val) = &value.kind { - sub.add_global(idx, val.clone()); - } else { - sub.globals.remove(&idx); - } - } - _ => {} - } - (BoundKind::Set { addr, value }, node.ty) - }, + BoundKind::Tuple { elements } => { let elements = elements.into_iter().map(|e| self.visit_node(e, sub)).collect(); (BoundKind::Tuple { elements }, node.ty) @@ -338,11 +365,24 @@ impl Optimizer { 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) => { - // Inline Closures ONLY if they are stateless (cracked) if let Some(closure) = obj.as_any().downcast_ref::() { closure.upvalues.is_empty() } else { @@ -365,7 +405,7 @@ impl Optimizer { false } } - _ => true, // Locals and Upvalues are pure in terms of side-effects + _ => true, } } BoundKind::Tuple { elements } => elements.iter().all(|e| self.is_pure(e)), @@ -376,9 +416,7 @@ impl Optimizer { BoundKind::Block { exprs } => exprs.iter().all(|e| self.is_pure(e)), BoundKind::Expansion { bound_expanded, .. } => self.is_pure(bound_expanded), BoundKind::DefLocal { value, .. } | BoundKind::Set { addr: Address::Local(_), value } => self.is_pure(value), - BoundKind::Call { callee, args } => { - // A call is pure if the callee is a known pure function and args are pure let callee_pure = match &callee.kind { BoundKind::Get { addr: Address::Global(idx), .. } => { if let Some(purity_rc) = &self.global_purity { @@ -387,7 +425,7 @@ impl Optimizer { false } } - BoundKind::Constant(Value::Function(_)) => true, // Native fns are checked at registry + BoundKind::Constant(Value::Function(_)) => true, BoundKind::Constant(Value::Object(obj)) => { if let Some(closure) = obj.as_any().downcast_ref::() { self.is_pure(&closure.function_node) @@ -399,38 +437,24 @@ impl Optimizer { }; callee_pure && self.is_pure(args) } - - _ => false, // DefGlobal and Set(Global) are impure + _ => false, } } fn try_fold_pure(&self, callee: &TypedNode, args: &TypedNode) -> Option { - // 1. Check if the whole call is pure - // Create a temporary call node to check purity let temp_call = Node { identity: callee.identity.clone(), ty: StaticType::Any, kind: BoundKind::Call { callee: Box::new(callee.clone()), args: Box::new(args.clone()) } }; - - if !self.is_pure(&temp_call) { - return None; - } - - // 2. Extract constant arguments + if !self.is_pure(&temp_call) { 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; // All arguments must be constants for folding - } + if let BoundKind::Constant(val) = node.kind { arg_values.push(val); } + else { return None; } } - - // 3. Resolve the function to execute let func_val = match &callee.kind { BoundKind::Get { addr: Address::Global(idx), .. } => { self.globals.as_ref()?.borrow().get(*idx as usize)?.clone() @@ -438,26 +462,10 @@ impl Optimizer { BoundKind::Constant(val) => val.clone(), _ => return None, }; - - // 4. Execute let result = match func_val { Value::Function(f) => f(arg_values), - Value::Object(obj) => { - if let Some(_closure) = obj.as_any().downcast_ref::() { - // For closures, we'd need a VM to execute. - // However, if it's pure and we are in the Optimizer, - // we might want to avoid full VM spin-up here if possible, - // or use a shared VM instance. - // For now, we only fold Native Functions. - return None; - } else { - return None; - } - } _ => return None, }; - - // 5. Return new constant node Some(Node { identity: callee.identity.clone(), ty: result.static_type(), @@ -466,10 +474,6 @@ impl Optimizer { } fn try_beta_reduce(&self, params: &TypedNode, args: &TypedNode, body: TypedNode) -> Option { - if self.contains_def_local(&body) { - return None; - } - let mut sub = SubstitutionMap::new(); let mut arg_vals = Vec::new(); self.flatten_typed_tuple(args, &mut arg_vals); @@ -477,22 +481,15 @@ impl Optimizer { let mut slot_index = 0; self.map_params_to_args(params, &arg_vals, &mut slot_index, &mut sub); - if sub.locals.len() < arg_vals.len() { - return None; - } - - if sub.locals.is_empty() && sub.upvalues.is_empty() && !arg_vals.is_empty() { - return None; - } + if sub.locals.len() < arg_vals.len() { return None; } + if sub.locals.is_empty() && sub.upvalues.is_empty() && !arg_vals.is_empty() { return None; } Some(self.visit_node(body, &mut sub)) } 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); - } + for el in elements { self.flatten_typed_tuple(el, into); } } else if !matches!(node.kind, BoundKind::Nop) { into.push(node.clone()); } @@ -505,47 +502,70 @@ impl Optimizer { && let BoundKind::Constant(val) = &arg.kind { sub.add_local(*slot, val.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); - } + for el in elements { self.map_params_to_args(el, args, offset, sub); } } _ => {} } } - fn contains_def_local(&self, node: &TypedNode) -> bool { + fn collect_local_usage(&self, node: &TypedNode, used: &mut HashSet) { match &node.kind { - BoundKind::DefLocal { .. } => true, - BoundKind::If { cond, then_br, else_br } => { - self.contains_def_local(cond) || self.contains_def_local(then_br) || else_br.as_ref().is_some_and(|e| self.contains_def_local(e)) + BoundKind::Get { addr: Address::Local(slot), .. } | BoundKind::Set { addr: Address::Local(slot), .. } => { + used.insert(*slot); } - BoundKind::Block { exprs } => exprs.iter().any(|e| self.contains_def_local(e)), - BoundKind::Call { callee, args } => self.contains_def_local(callee) || self.contains_def_local(args), - BoundKind::Lambda { .. } => false, - BoundKind::Tuple { elements } => elements.iter().any(|e| self.contains_def_local(e)), - BoundKind::Record { fields } => fields.iter().any(|(k, v)| self.contains_def_local(k) || self.contains_def_local(v)), - BoundKind::Set { value, .. } => self.contains_def_local(value), - _ => false, + BoundKind::Lambda { upvalues, .. } => { + for addr in upvalues { + if let Address::Local(slot) = addr { used.insert(*slot); } + } + } + BoundKind::Block { exprs } => { + for e in exprs { self.collect_local_usage(e, used); } + } + BoundKind::If { cond, then_br, else_br } => { + self.collect_local_usage(cond, used); + self.collect_local_usage(then_br, used); + if let Some(e) = else_br { self.collect_local_usage(e, used); } + } + BoundKind::Call { callee, args } => { + self.collect_local_usage(callee, used); + self.collect_local_usage(args, used); + } + BoundKind::Tuple { elements } => { + for e in elements { self.collect_local_usage(e, used); } + } + BoundKind::Record { fields } => { + for (k, v) in fields { + self.collect_local_usage(k, used); + self.collect_local_usage(v, used); + } + } + BoundKind::DefLocal { value, .. } | BoundKind::Set { value, .. } => { + self.collect_local_usage(value, used); + } + BoundKind::Expansion { bound_expanded, .. } => { + self.collect_local_usage(bound_expanded, used); + } + _ => {} } } } /// Helper for transitively inlining values and cleaning up capture lists. struct SubstitutionMap { - /// Mapping Address::Local(slot) -> Value locals: HashMap, - /// Mapping Address::Upvalue(idx) -> Value upvalues: HashMap, - /// Mapping Address::Global(idx) -> Value (Script-local substitution) globals: HashMap, - /// Tracking which local slots are actually used (not inlined) + slot_mapping: HashMap, + next_slot: u32, used_slots: HashSet, - /// Tracking which global indices are used + pure_slots: HashSet, used_globals: HashSet, - /// Tracking which local slots are captured by lambdas captured_slots: HashSet, } @@ -555,12 +575,25 @@ impl SubstitutionMap { locals: HashMap::new(), upvalues: HashMap::new(), globals: HashMap::new(), + slot_mapping: HashMap::new(), + next_slot: 0, used_slots: HashSet::new(), + pure_slots: HashSet::new(), used_globals: HashSet::new(), captured_slots: HashSet::new(), } } + 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); } @@ -573,7 +606,6 @@ impl SubstitutionMap { self.globals.insert(idx, val); } - /// Re-maps Get(Upvalue(old_idx)) to Get(Upvalue(new_idx)) based on a mapping table. fn reindex_upvalues(&self, node: TypedNode, mapping: &[Option]) -> TypedNode { let (new_kind, new_ty) = match node.kind { BoundKind::Get { addr: Address::Upvalue(idx), name } => { @@ -586,22 +618,18 @@ impl SubstitutionMap { (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)); - } + 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)); @@ -639,7 +667,6 @@ impl SubstitutionMap { }, k => (k, node.ty), }; - Node { identity: node.identity, kind: new_kind, ty: new_ty } } } @@ -662,19 +689,12 @@ mod tests { #[test] fn test_opt_folding_complex() { - // Test float folding let dump_float = get_optimized_dump("(+ 1.5 2.5)", 2); assert!(dump_float.contains("Constant: 4")); - - // Test text folding let dump_text = get_optimized_dump("(+ \"hello \" \"world\")", 2); assert!(dump_text.contains("Constant: \"hello world\"")); - - // Test date folding let dump_date = get_optimized_dump("(date \"2023-01-01\")", 2); assert!(dump_date.contains("Constant: #2023-01-01 00:00:00#")); - - // Test comparison folding let dump_cmp = get_optimized_dump("(> 10 5)", 2); assert!(dump_cmp.contains("Constant: true")); } diff --git a/src/ast/environment.rs b/src/ast/environment.rs index 53b7fa7..b898f86 100644 --- a/src/ast/environment.rs +++ b/src/ast/environment.rs @@ -91,7 +91,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, + optimization_level: 2, }; env.register_stdlib(); env