From 3142e64bbdf5f4ae513b9211ba806f0c7f0275e0 Mon Sep 17 00:00:00 2001 From: Michael Schimmel Date: Sun, 22 Feb 2026 10:41:48 +0100 Subject: [PATCH] Refactor recursion checking Introduces a `UsageInfo` struct to consolidate tracking of used and assigned locals, globals, and identities. Replaces multiple `HashSet` arguments in `collect_usage` with a single `UsageInfo` struct. Adds an `is_recursive` method to check for recursive calls within a node, considering global indices, local slots, and identity. Updates inlining logic to use `is_recursive` to prevent inlining of recursive functions or closures. This includes: - Checking for recursion when inlining a call to a lambda. - Checking for recursion when inlining a global variable that is a lambda. - Checking for recursion when inlining a closure. Updates `is_inlinable_value` to also check for closure recursion. --- src/ast/compiler/optimizer.rs | 341 +++++++----------- ...optimizer__tests__opt_lambda_cracking.snap | 7 +- 2 files changed, 135 insertions(+), 213 deletions(-) diff --git a/src/ast/compiler/optimizer.rs b/src/ast/compiler/optimizer.rs index 3915aa2..61a4724 100644 --- a/src/ast/compiler/optimizer.rs +++ b/src/ast/compiler/optimizer.rs @@ -54,6 +54,15 @@ impl PathTracker { } } +#[derive(Default)] +struct UsageInfo { + used_locals: HashSet, + used_globals: HashSet, + used_identities: HashSet, + assigned_locals: HashSet, + assigned_globals: HashSet, +} + impl Optimizer { pub fn new(enabled: bool) -> Self { Self { @@ -98,6 +107,34 @@ impl Optimizer { current } + fn is_recursive( + &self, + node: &TypedNode, + global_idx: Option, + local_slot: Option, + identity: Option<&crate::ast::types::Identity>, + ) -> bool { + let mut info = UsageInfo::default(); + self.collect_usage(node, &mut info); + + if let Some(idx) = global_idx { + if info.used_globals.contains(&idx) { + return true; + } + } + if let Some(slot) = local_slot { + if info.used_locals.contains(&slot) { + return true; + } + } + if let Some(id) = identity { + if info.used_identities.contains(id) { + return true; + } + } + false + } + fn visit_node( &self, node: TypedNode, @@ -136,16 +173,18 @@ impl Optimizer { 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 self.is_inlinable_value(val, Some(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) + && self.is_inlinable_value(val, Some(idx)) { return Node { identity: node.identity, @@ -225,7 +264,10 @@ impl Optimizer { && positional_count.is_some() && path.inlining_depth < 5 { - if path.enter_lambda(&callee.identity) { + // STRICT RECURSION CHECK: Don't inline if recursive + if !self.is_recursive(body, None, None, Some(&callee.identity)) + && path.enter_lambda(&callee.identity) + { path.inlining_depth += 1; let collapsed = self.try_beta_reduce_with_sub( params, @@ -263,20 +305,8 @@ impl Optimizer { && 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) { + // STRICT RECURSION CHECK: Don't inline if recursive (idx or identity) + if !self.is_recursive(body, Some(*idx), None, Some(&lambda_node.identity)) { let mut inner_sub = SubstitutionMap::new(); path.inlining_stack.insert(*idx); path.inlining_depth += 1; @@ -304,29 +334,39 @@ impl Optimizer { && (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()); - } + // STRICT RECURSION CHECK: Don't inline if closure body contains its own identity + if !self.is_recursive( + &closure.function_node, + None, + None, + Some(&closure.function_node.identity), + ) && path.enter_lambda(&closure.function_node.identity) + { + 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 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; + 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(res) = collapsed { + return res; + } } } @@ -342,6 +382,7 @@ impl Optimizer { && (closure.upvalues.is_empty() || self.purity_of(&closure.function_node) >= Purity::SideEffectFree) { + // CRACKING fallback (only if not recursive already handled above) let mut closure_sub = SubstitutionMap::new(); for (i, cell) in closure.upvalues.iter().enumerate() { closure_sub.add_upvalue(i as u32, cell.borrow().clone()); @@ -415,20 +456,11 @@ impl Optimizer { } 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(); + let mut info = UsageInfo::default(); 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, - ); + self.collect_usage(e, &mut info); } if let Some(last) = exprs.last() { match &last.kind { @@ -436,13 +468,13 @@ impl Optimizer { addr: Address::Local(slot), .. } => { - used_in_block.insert(*slot); + info.used_locals.insert(*slot); } BoundKind::Get { addr: Address::Global(idx), .. } => { - used_globals_in_block.insert(*idx); + info.used_globals.insert(*idx); } _ => {} } @@ -451,10 +483,10 @@ impl Optimizer { // 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 { + for slot in &info.assigned_locals { sub.assigned_locals.insert(*slot); } - for idx in &assigned_globals_in_block { + for idx in &info.assigned_globals { sub.assigned_globals.insert(*idx); } @@ -467,7 +499,7 @@ impl Optimizer { if self.enabled && !is_last { let removable = match &e.kind { BoundKind::DefLocal { slot, value, .. } => { - !used_in_block.contains(slot) + !info.used_locals.contains(slot) && !sub.captured_slots.contains(slot) && self.purity_of(value) >= Purity::SideEffectFree } @@ -476,7 +508,7 @@ impl Optimizer { value, .. } => { - !used_globals_in_block.contains(global_index) + !info.used_globals.contains(global_index) && (self.purity_of(value) >= Purity::SideEffectFree || matches!(value.kind, BoundKind::Lambda { .. })) } @@ -485,7 +517,7 @@ impl Optimizer { value, .. } => { - !used_in_block.contains(slot) + !info.used_locals.contains(slot) && !sub.captured_slots.contains(slot) && self.purity_of(value) >= Purity::SideEffectFree } @@ -538,29 +570,19 @@ impl Optimizer { 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, - ); + let mut info = UsageInfo::default(); + self.collect_usage(&node, &mut info); // Track captures from outer scope that are assigned in this lambda - for idx in &assigned_globals_in_lambda { + for idx in &info.assigned_globals { 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; + next_inner_subs.assigned_locals = info.assigned_locals; + next_inner_subs.assigned_globals = info.assigned_globals; for (old_idx, capture_addr) in original_upvalues.iter().enumerate() { let mut inlined_val = None; @@ -660,7 +682,9 @@ impl Optimizer { } => { let value = Box::new(self.visit_node(*value, sub, path)); if let BoundKind::Constant(val) = &value.kind { - sub.add_global(global_index, val.clone()); + if self.is_inlinable_value(val, Some(global_index)) { + sub.add_global(global_index, val.clone()); + } } let p = self.purity_of(&value); if p > Purity::Impure @@ -738,7 +762,7 @@ impl Optimizer { } } - fn is_inlinable_value(&self, val: &Value) -> bool { + fn is_inlinable_value(&self, val: &Value, global_idx: Option) -> bool { match val { Value::Int(_) | Value::Float(_) @@ -748,7 +772,14 @@ impl Optimizer { | Value::DateTime(_) => true, Value::Object(obj) => { if let Some(closure) = obj.as_any().downcast_ref::() { + // A closure is only inlinable if it's not recursive closure.upvalues.is_empty() + && !self.is_recursive( + &closure.function_node, + global_idx, + None, + Some(&closure.function_node.identity), + ) } else { false } @@ -951,96 +982,46 @@ impl Optimizer { } } - fn collect_usage( - &self, - node: &TypedNode, - used_locals: &mut HashSet, - used_globals: &mut HashSet, - assigned_locals: &mut HashSet, - assigned_globals: &mut HashSet, - ) { + fn collect_usage(&self, node: &TypedNode, info: &mut UsageInfo) { 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, - ); + info.used_identities + .insert(closure.function_node.identity.clone()); + self.collect_usage(&closure.function_node, info); } } BoundKind::Get { addr, .. } => match addr { Address::Local(slot) => { - used_locals.insert(*slot); + info.used_locals.insert(*slot); } Address::Global(idx) => { - used_globals.insert(*idx); + info.used_globals.insert(*idx); } _ => {} }, BoundKind::Set { addr, value } => { match addr { Address::Local(slot) => { - assigned_locals.insert(*slot); + info.assigned_locals.insert(*slot); } Address::Global(idx) => { - assigned_globals.insert(*idx); + info.assigned_globals.insert(*idx); } _ => {} } - self.collect_usage( - value, - used_locals, - used_globals, - assigned_locals, - assigned_globals, - ); + self.collect_usage(value, info); } - 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::Lambda { params, body, .. } => { + info.used_identities.insert(node.identity.clone()); + self.collect_usage(params, info); + self.collect_usage(body, info); } BoundKind::Block { exprs } => { for e in exprs { - self.collect_usage( - e, - used_locals, - used_globals, - assigned_locals, - assigned_globals, - ); + self.collect_usage(e, info); } } BoundKind::If { @@ -1048,92 +1029,32 @@ impl Optimizer { 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, - ); + self.collect_usage(cond, info); + self.collect_usage(then_br, info); if let Some(e) = else_br { - self.collect_usage( - e, - used_locals, - used_globals, - assigned_locals, - assigned_globals, - ); + self.collect_usage(e, info); } } 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, - ); + self.collect_usage(callee, info); + self.collect_usage(args, info); } BoundKind::Tuple { elements } => { for e in elements { - self.collect_usage( - e, - used_locals, - used_globals, - assigned_locals, - assigned_globals, - ); + self.collect_usage(e, info); } } 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, - ); + self.collect_usage(k, info); + self.collect_usage(v, info); } } BoundKind::DefLocal { value, .. } | BoundKind::DefGlobal { value, .. } => { - self.collect_usage( - value, - used_locals, - used_globals, - assigned_locals, - assigned_globals, - ); + self.collect_usage(value, info); } BoundKind::Expansion { bound_expanded, .. } => { - self.collect_usage( - bound_expanded, - used_locals, - used_globals, - assigned_locals, - assigned_globals, - ); + self.collect_usage(bound_expanded, info); } _ => {} } diff --git a/src/ast/compiler/snapshots/myc__ast__compiler__optimizer__tests__opt_lambda_cracking.snap b/src/ast/compiler/snapshots/myc__ast__compiler__optimizer__tests__opt_lambda_cracking.snap index 5e6fbbb..2fe6461 100644 --- a/src/ast/compiler/snapshots/myc__ast__compiler__optimizer__tests__opt_lambda_cracking.snap +++ b/src/ast/compiler/snapshots/myc__ast__compiler__optimizer__tests__opt_lambda_cracking.snap @@ -1,12 +1,13 @@ --- source: src/ast/compiler/optimizer.rs -assertion_line: 727 +assertion_line: 1333 expression: dump --- Lambda (Upvalues: 0) Parameters: Tuple - Lambda (Upvalues: 0) + Lambda (Upvalues: 1) Parameters: Tuple - Constant: 10 + Upvalues: [Local(0)] + Get: x (Upvalue(0))