From bd506b1b1795824e535538677c766eacf2729e1f Mon Sep 17 00:00:00 2001 From: Michael Schimmel Date: Sun, 22 Feb 2026 10:05:56 +0100 Subject: [PATCH] feat: Track optimization path to prevent infinite recursion Introduce a `PathTracker` to manage recursion depth and detect cycles during optimization. This prevents infinite inlining of recursive functions. --- src/ast/compiler/optimizer.rs | 190 ++++++++++++++++++++++------------ 1 file changed, 122 insertions(+), 68 deletions(-) diff --git a/src/ast/compiler/optimizer.rs b/src/ast/compiler/optimizer.rs index 7c008d6..3915aa2 100644 --- a/src/ast/compiler/optimizer.rs +++ b/src/ast/compiler/optimizer.rs @@ -23,6 +23,37 @@ pub struct Optimizer { 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 { @@ -57,7 +88,8 @@ impl Optimizer { let mut current = node; for _ in 0..self.max_passes { let mut sub = SubstitutionMap::new(); - let next = self.visit_node(current.clone(), &mut sub); + let mut path = PathTracker::new(); + let next = self.visit_node(current.clone(), &mut sub, &mut path); if next == current { break; } @@ -66,7 +98,12 @@ impl Optimizer { current } - fn visit_node(&self, node: TypedNode, sub: &mut SubstitutionMap) -> TypedNode { + 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 { @@ -143,7 +180,7 @@ impl Optimizer { } BoundKind::Set { addr, value } => { - let value = Box::new(self.visit_node(*value, sub)); + 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 { @@ -173,8 +210,8 @@ impl Optimizer { } BoundKind::Call { callee, args } => { - let callee = self.visit_node(*callee, sub); - let args = self.visit_node(*args, sub); + 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 @@ -186,10 +223,24 @@ impl Optimizer { } = &callee.kind && upvalues.is_empty() && positional_count.is_some() - && let Some(collapsed) = - self.try_beta_reduce(params, &args, (**body).clone()) + && path.inlining_depth < 5 { - return collapsed; + 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) @@ -198,8 +249,8 @@ impl Optimizer { .. } = &callee.kind && let Some(registry_rc) = &self.lambda_registry - && sub.inlining_depth < 5 - && !sub.inlining_stack.contains(idx) + && path.inlining_depth < 5 + && !path.inlining_stack.contains(idx) { let registry = registry_rc.borrow(); if let Some(lambda_node) = registry.get(idx) @@ -227,16 +278,20 @@ impl Optimizer { if !used_globals_in_body.contains(idx) { let mut inner_sub = SubstitutionMap::new(); - inner_sub.inlining_stack = sub.inlining_stack.clone(); - inner_sub.inlining_stack.insert(*idx); - inner_sub.inlining_depth = sub.inlining_depth + 1; - if let Some(collapsed) = self.try_beta_reduce_with_sub( + 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, - ) { - return collapsed; + path, + ); + path.inlining_depth -= 1; + path.inlining_stack.remove(idx); + + if let Some(res) = collapsed { + return res; } } } @@ -244,26 +299,34 @@ impl Optimizer { // Case 2: Cracking and Inlining for Constant Closures if let BoundKind::Constant(Value::Object(ref obj)) = callee.kind - && sub.inlining_depth < 5 + && 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(); - closure_sub.inlining_depth = sub.inlining_depth + 1; for (i, cell) in closure.upvalues.iter().enumerate() { closure_sub.add_upvalue(i as u32, cell.borrow().clone()); } - let inlined_body = - self.visit_node((*closure.function_node).clone(), &mut closure_sub); - if let Some(collapsed) = self.try_beta_reduce_with_sub( + 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, - ) { - return collapsed; + path, + ); + path.inlining_depth -= 1; + + if let Some(res) = collapsed { + return res; } } @@ -274,18 +337,19 @@ impl Optimizer { if self.enabled && let BoundKind::Constant(Value::Object(ref obj)) = callee.kind - && sub.inlining_depth < 5 + && 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(); - closure_sub.inlining_depth = sub.inlining_depth + 1; 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); + self.visit_node((*closure.function_node).clone(), &mut closure_sub, path); + path.inlining_depth -= 1; let cracked_lambda = Node { identity: callee.identity.clone(), @@ -321,14 +385,14 @@ impl Optimizer { then_br, else_br, } => { - let cond = self.visit_node(*cond, sub); + 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); + return self.visit_node(*then_br, sub, path); } else if let Some(else_node) = else_br { - return self.visit_node(*else_node, sub); + return self.visit_node(*else_node, sub, path); } else { return Node { identity: node.identity, @@ -338,8 +402,8 @@ impl Optimizer { } } - let then_br = Box::new(self.visit_node(*then_br, sub)); - let else_br = else_br.map(|e| Box::new(self.visit_node(*e, sub))); + 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), @@ -433,7 +497,7 @@ impl Optimizer { } } - let opt = self.visit_node(e, sub); + let opt = self.visit_node(e, sub, path); if self.enabled && matches!(opt.kind, BoundKind::Nop) && !is_last { continue; } @@ -495,7 +559,6 @@ impl Optimizer { let mut new_upvalues = Vec::new(); let mut mapping = Vec::new(); let mut next_inner_subs = SubstitutionMap::new(); - next_inner_subs.inlining_depth = sub.inlining_depth; next_inner_subs.assigned_locals = assigned_locals_in_lambda; next_inner_subs.assigned_globals = assigned_globals_in_lambda; @@ -531,13 +594,14 @@ impl Optimizer { } // 1. Visit parameters to determine their new slots - let params_node = self.visit_node(params.as_ref().clone(), &mut next_inner_subs); + 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); + 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 { @@ -564,7 +628,7 @@ impl Optimizer { if !captured_by.is_empty() { sub.captured_slots.insert(slot); } - let value = Box::new(self.visit_node(*value, sub)); + let value = Box::new(self.visit_node(*value, sub, path)); if let BoundKind::Constant(val) = &value.kind { sub.add_local(slot, val.clone()); } else { @@ -594,7 +658,7 @@ impl Optimizer { global_index, value, } => { - let value = Box::new(self.visit_node(*value, sub)); + let value = Box::new(self.visit_node(*value, sub, path)); if let BoundKind::Constant(val) = &value.kind { sub.add_global(global_index, val.clone()); } @@ -617,14 +681,19 @@ impl Optimizer { BoundKind::Tuple { elements } => { let elements = elements .into_iter() - .map(|e| self.visit_node(e, sub)) + .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), self.visit_node(v, sub))) + .map(|(k, v)| { + ( + self.visit_node(k, sub, path), + self.visit_node(v, sub, path), + ) + }) .collect(); (BoundKind::Record { fields }, node.ty) } @@ -632,9 +701,9 @@ impl Optimizer { original_call, bound_expanded, } => { - sub.inlining_depth += 1; - let bound_expanded = Box::new(self.visit_node(*bound_expanded, sub)); - sub.inlining_depth -= 1; + path.inlining_depth += 1; + let bound_expanded = Box::new(self.visit_node(*bound_expanded, sub, path)); + path.inlining_depth -= 1; ( BoundKind::Expansion { original_call, @@ -804,22 +873,13 @@ impl Optimizer { }) } - fn try_beta_reduce( - &self, - params: &TypedNode, - args: &TypedNode, - body: TypedNode, - ) -> Option { - let mut sub = SubstitutionMap::new(); - self.try_beta_reduce_with_sub(params, args, body, &mut sub) - } - 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); @@ -840,7 +900,7 @@ impl Optimizer { return None; } - Some(self.visit_node(body, sub)) + Some(self.visit_node(body, sub, path)) } fn flatten_typed_tuple(&self, node: &TypedNode, into: &mut Vec) { @@ -1087,8 +1147,6 @@ struct SubstitutionMap { globals: HashMap, ast_locals: HashMap, slot_mapping: HashMap, - inlining_stack: HashSet, - inlining_depth: usize, assigned_locals: HashSet, assigned_globals: HashSet, next_slot: u32, @@ -1106,8 +1164,6 @@ impl SubstitutionMap { globals: HashMap::new(), ast_locals: HashMap::new(), slot_mapping: HashMap::new(), - inlining_stack: HashSet::new(), - inlining_depth: 0, assigned_locals: HashSet::new(), assigned_globals: HashSet::new(), next_slot: 0, @@ -1369,17 +1425,15 @@ mod tests { } #[test] - fn test_opt_unused_lambda_with_side_effects_dce() { - // Die Lambda-Definition selbst ist Pure, auch wenn ihr Body Impure ist (assign). - // Daher muss sie gelöscht werden, wenn sie nicht gerufen wird. - // 'g' muss definiert sein, damit der Binder nicht abbricht. - let source = "(fn [] (do (def g 0) (def f (fn [x] (assign g 1))) 42))"; + 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); - assert!( - !dump.contains("DefLocal (Name: 'f'"), - "Unused lambda with side effects should be removed by DCE. Dump: \n{}", - dump - ); - assert!(dump.contains("Constant: 42")); + + // 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); } }