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>>>>,
}
/// 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 {
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>()
&& (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>()
&& (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(&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() {
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<TypedNode> {
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<TypedNode> {
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<TypedNode>) {
@@ -1087,8 +1147,6 @@ struct SubstitutionMap {
globals: HashMap<u32, Value>,
ast_locals: HashMap<u32, TypedNode>,
slot_mapping: HashMap<u32, u32>,
inlining_stack: HashSet<u32>,
inlining_depth: usize,
assigned_locals: HashSet<u32>,
assigned_globals: HashSet<u32>,
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);
}
}