Brummel/RustAst
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Track optimization path to prevent infinite recursion

Browse Source 
Introduce a `PathTracker` to manage recursion depth and detect cycles
during optimization. This prevents infinite inlining of recursive
functions.
This commit is contained in:
Michael Schimmel
2026-02-22 10:05:56 +01:00
parent e30444e89b
commit bd506b1b17
1 changed files with 122 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ast/compiler/optimizer.rs
+122 -68
View File
@@ -23,6 +23,37 @@ pub struct Optimizer {
pub lambda_registry: Option<Rc<RefCell<HashMap<u32, TypedNode>>>>, pub lambda_registry: Option<Rc<RefCell<HashMap<u32, TypedNode>>>>,
} }
/// 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<u32>,
/// Stack of node identities to detect recursion in anonymous lambdas.
identity_stack: HashSet<crate::ast::types::Identity>,
}
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 { impl Optimizer {
pub fn new(enabled: bool) -> Self { pub fn new(enabled: bool) -> Self {
Self { Self {
@@ -57,7 +88,8 @@ impl Optimizer {
let mut current = node; let mut current = node;
for _ in 0..self.max_passes { for _ in 0..self.max_passes {
let mut sub = SubstitutionMap::new(); 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 { if next == current {
break; break;
} }
@@ -66,7 +98,12 @@ impl Optimizer {
current 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 { let (new_kind, new_ty) = match node.kind {
BoundKind::Get { addr, name } => { BoundKind::Get { addr, name } => {
let new_addr = match addr { let new_addr = match addr {
@@ -143,7 +180,7 @@ impl Optimizer {
} }
BoundKind::Set { addr, value } => { 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 { let new_addr = match addr {
Address::Local(slot) => { Address::Local(slot) => {
if let BoundKind::Constant(val) = &value.kind { if let BoundKind::Constant(val) = &value.kind {
@@ -173,8 +210,8 @@ impl Optimizer {
} }
BoundKind::Call { callee, args } => { BoundKind::Call { callee, args } => {
let callee = self.visit_node(*callee, sub); let callee = self.visit_node(*callee, sub, path);
let args = self.visit_node(*args, sub); let args = self.visit_node(*args, sub, path);
if self.enabled { if self.enabled {
// Case 1: Beta-Reduction for Lambda Literals // Case 1: Beta-Reduction for Lambda Literals
@@ -186,10 +223,24 @@ impl Optimizer {
} = &callee.kind } = &callee.kind
&& upvalues.is_empty() && upvalues.is_empty()
&& positional_count.is_some() && positional_count.is_some()
&& let Some(collapsed) = && path.inlining_depth < 5
self.try_beta_reduce(params, &args, (**body).clone())
{ {
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) // Case 1.5: Beta-Reduction for Global Registry functions (from Registry)
@@ -198,8 +249,8 @@ impl Optimizer {
.. ..
} = &callee.kind } = &callee.kind
&& let Some(registry_rc) = &self.lambda_registry && let Some(registry_rc) = &self.lambda_registry
&& sub.inlining_depth < 5 && path.inlining_depth < 5
&& !sub.inlining_stack.contains(idx) && !path.inlining_stack.contains(idx)
{ {
let registry = registry_rc.borrow(); let registry = registry_rc.borrow();
if let Some(lambda_node) = registry.get(idx) if let Some(lambda_node) = registry.get(idx)
@@ -227,16 +278,20 @@ impl Optimizer {
if !used_globals_in_body.contains(idx) { if !used_globals_in_body.contains(idx) {
let mut inner_sub = SubstitutionMap::new(); let mut inner_sub = SubstitutionMap::new();
inner_sub.inlining_stack = sub.inlining_stack.clone(); path.inlining_stack.insert(*idx);
inner_sub.inlining_stack.insert(*idx); path.inlining_depth += 1;
inner_sub.inlining_depth = sub.inlining_depth + 1; let collapsed = self.try_beta_reduce_with_sub(
if let Some(collapsed) = self.try_beta_reduce_with_sub(
params.as_ref(), params.as_ref(),
&args, &args,
body.as_ref().clone(), body.as_ref().clone(),
&mut inner_sub, &mut inner_sub,
) { path,
return collapsed; );
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 // Case 2: Cracking and Inlining for Constant Closures
if let BoundKind::Constant(Value::Object(ref obj)) = callee.kind 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>() && let Some(closure) = obj.as_any().downcast_ref::<Closure>()
&& (closure.upvalues.is_empty() && (closure.upvalues.is_empty()
|| self.purity_of(&closure.function_node) >= Purity::SideEffectFree) || self.purity_of(&closure.function_node) >= Purity::SideEffectFree)
{ {
let mut closure_sub = SubstitutionMap::new(); let mut closure_sub = SubstitutionMap::new();
closure_sub.inlining_depth = sub.inlining_depth + 1;
for (i, cell) in closure.upvalues.iter().enumerate() { for (i, cell) in closure.upvalues.iter().enumerate() {
closure_sub.add_upvalue(i as u32, cell.borrow().clone()); 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, &closure.parameter_node,
&args, &args,
inlined_body, inlined_body,
&mut closure_sub, &mut closure_sub,
) { path,
return collapsed; );
path.inlining_depth -= 1;
if let Some(res) = collapsed {
return res;
} }
} }
@@ -274,18 +337,19 @@ impl Optimizer {
if self.enabled if self.enabled
&& let BoundKind::Constant(Value::Object(ref obj)) = callee.kind && 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>() && let Some(closure) = obj.as_any().downcast_ref::<Closure>()
&& (closure.upvalues.is_empty() && (closure.upvalues.is_empty()
|| self.purity_of(&closure.function_node) >= Purity::SideEffectFree) || self.purity_of(&closure.function_node) >= Purity::SideEffectFree)
{ {
let mut closure_sub = SubstitutionMap::new(); let mut closure_sub = SubstitutionMap::new();
closure_sub.inlining_depth = sub.inlining_depth + 1;
for (i, cell) in closure.upvalues.iter().enumerate() { for (i, cell) in closure.upvalues.iter().enumerate() {
closure_sub.add_upvalue(i as u32, cell.borrow().clone()); closure_sub.add_upvalue(i as u32, cell.borrow().clone());
} }
path.inlining_depth += 1;
let inlined_body = 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 { let cracked_lambda = Node {
identity: callee.identity.clone(), identity: callee.identity.clone(),
@@ -321,14 +385,14 @@ impl Optimizer {
then_br, then_br,
else_br, else_br,
} => { } => {
let cond = self.visit_node(*cond, sub); let cond = self.visit_node(*cond, sub, path);
if self.enabled if self.enabled
&& let BoundKind::Constant(ref val) = cond.kind && let BoundKind::Constant(ref val) = cond.kind
{ {
if val.is_truthy() { 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 { } else if let Some(else_node) = else_br {
return self.visit_node(*else_node, sub); return self.visit_node(*else_node, sub, path);
} else { } else {
return Node { return Node {
identity: node.identity, identity: node.identity,
@@ -338,8 +402,8 @@ impl Optimizer {
} }
} }
let then_br = Box::new(self.visit_node(*then_br, 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))); let else_br = else_br.map(|e| Box::new(self.visit_node(*e, sub, path)));
( (
BoundKind::If { BoundKind::If {
cond: Box::new(cond), 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 { if self.enabled && matches!(opt.kind, BoundKind::Nop) && !is_last {
continue; continue;
} }
@@ -495,7 +559,6 @@ impl Optimizer {
let mut new_upvalues = Vec::new(); let mut new_upvalues = Vec::new();
let mut mapping = Vec::new(); let mut mapping = Vec::new();
let mut next_inner_subs = SubstitutionMap::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_locals = assigned_locals_in_lambda;
next_inner_subs.assigned_globals = assigned_globals_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 // 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 // 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. // so they are mapped 1:1 and don't get reassigned to higher slots in the body.
self.collect_parameter_slots(&params_node, &mut next_inner_subs); self.collect_parameter_slots(&params_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() { let reindexed_body = if new_upvalues.len() != original_upvalues.len() {
sub.reindex_upvalues(body_node, &mapping) sub.reindex_upvalues(body_node, &mapping)
} else { } else {
@@ -564,7 +628,7 @@ impl Optimizer {
if !captured_by.is_empty() { if !captured_by.is_empty() {
sub.captured_slots.insert(slot); 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 { if let BoundKind::Constant(val) = &value.kind {
sub.add_local(slot, val.clone()); sub.add_local(slot, val.clone());
} else { } else {
@@ -594,7 +658,7 @@ impl Optimizer {
global_index, global_index,
value, 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 { if let BoundKind::Constant(val) = &value.kind {
sub.add_global(global_index, val.clone()); sub.add_global(global_index, val.clone());
} }
@@ -617,14 +681,19 @@ impl Optimizer {
BoundKind::Tuple { elements } => { BoundKind::Tuple { elements } => {
let elements = elements let elements = elements
.into_iter() .into_iter()
.map(|e| self.visit_node(e, sub)) .map(|e| self.visit_node(e, sub, path))
.collect(); .collect();
(BoundKind::Tuple { elements }, node.ty) (BoundKind::Tuple { elements }, node.ty)
} }
BoundKind::Record { fields } => { BoundKind::Record { fields } => {
let fields = fields let fields = fields
.into_iter() .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(); .collect();
(BoundKind::Record { fields }, node.ty) (BoundKind::Record { fields }, node.ty)
} }
@@ -632,9 +701,9 @@ impl Optimizer {
original_call, original_call,
bound_expanded, bound_expanded,
} => { } => {
sub.inlining_depth += 1; path.inlining_depth += 1;
let bound_expanded = Box::new(self.visit_node(*bound_expanded, sub)); let bound_expanded = Box::new(self.visit_node(*bound_expanded, sub, path));
sub.inlining_depth -= 1; path.inlining_depth -= 1;
( (
BoundKind::Expansion { BoundKind::Expansion {
original_call, original_call,
@@ -804,22 +873,13 @@ impl Optimizer {
}) })
} }
fn try_beta_reduce(
&self,
params: &TypedNode,
args: &TypedNode,
body: TypedNode,
) -> Option<TypedNode> {
let mut sub = SubstitutionMap::new();
self.try_beta_reduce_with_sub(params, args, body, &mut sub)
}
fn try_beta_reduce_with_sub( fn try_beta_reduce_with_sub(
&self, &self,
params: &TypedNode, params: &TypedNode,
args: &TypedNode, args: &TypedNode,
body: TypedNode, body: TypedNode,
sub: &mut SubstitutionMap, sub: &mut SubstitutionMap,
path: &mut PathTracker,
) -> Option<TypedNode> { ) -> Option<TypedNode> {
let mut arg_vals = Vec::new(); let mut arg_vals = Vec::new();
self.flatten_typed_tuple(args, &mut arg_vals); self.flatten_typed_tuple(args, &mut arg_vals);
@@ -840,7 +900,7 @@ impl Optimizer {
return None; 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<TypedNode>) { fn flatten_typed_tuple(&self, node: &TypedNode, into: &mut Vec<TypedNode>) {
@@ -1087,8 +1147,6 @@ struct SubstitutionMap {
globals: HashMap<u32, Value>, globals: HashMap<u32, Value>,
ast_locals: HashMap<u32, TypedNode>, ast_locals: HashMap<u32, TypedNode>,
slot_mapping: HashMap<u32, u32>, slot_mapping: HashMap<u32, u32>,
inlining_stack: HashSet<u32>,
inlining_depth: usize,
assigned_locals: HashSet<u32>, assigned_locals: HashSet<u32>,
assigned_globals: HashSet<u32>, assigned_globals: HashSet<u32>,
next_slot: u32, next_slot: u32,
@@ -1106,8 +1164,6 @@ impl SubstitutionMap {
globals: HashMap::new(), globals: HashMap::new(),
ast_locals: HashMap::new(), ast_locals: HashMap::new(),
slot_mapping: HashMap::new(), slot_mapping: HashMap::new(),
inlining_stack: HashSet::new(),
inlining_depth: 0,
assigned_locals: HashSet::new(), assigned_locals: HashSet::new(),
assigned_globals: HashSet::new(), assigned_globals: HashSet::new(),
next_slot: 0, next_slot: 0,
@@ -1369,17 +1425,15 @@ mod tests {
} }
#[test] #[test]
fn test_opt_unused_lambda_with_side_effects_dce() { fn test_opt_local_recursion_safety() {
// Die Lambda-Definition selbst ist Pure, auch wenn ihr Body Impure ist (assign). // (fn [] (do (def f (fn [x] (f x))) (f 1)))
// Daher muss sie gelöscht werden, wenn sie nicht gerufen wird. // Hier sollte der Optimizer 'f' nicht unendlich oft in sich selbst inlinen.
// 'g' muss definiert sein, damit der Binder nicht abbricht. let source = "(fn [] (do (def f (fn [x] (f x))) (f 1)))";
let source = "(fn [] (do (def g 0) (def f (fn [x] (assign g 1))) 42))";
let dump = get_optimized_dump(source, true); let dump = get_optimized_dump(source, true);
assert!(
!dump.contains("DefLocal (Name: 'f'"), // Der Dump sollte 'Call' auf 'f' (oder Slot-Get) noch enthalten,
"Unused lambda with side effects should be removed by DCE. Dump: \n{}", // statt 5 Ebenen tief entfaltet zu sein.
dump 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);
assert!(dump.contains("Constant: 42"));
} }
} }