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

Refactor recursion checking

Browse Source 
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.
This commit is contained in:
Michael Schimmel
2026-02-22 10:41:48 +01:00
parent bd506b1b17
commit 3142e64bbd
2 changed files with 135 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ast/compiler/optimizer.rs
+131 -210
View File
@@ -54,6 +54,15 @@ impl PathTracker {
} }
} }
#[derive(Default)]
struct UsageInfo {
used_locals: HashSet<u32>,
used_globals: HashSet<u32>,
used_identities: HashSet<crate::ast::types::Identity>,
assigned_locals: HashSet<u32>,
assigned_globals: HashSet<u32>,
}
impl Optimizer { impl Optimizer {
pub fn new(enabled: bool) -> Self { pub fn new(enabled: bool) -> Self {
Self { Self {
@@ -98,6 +107,34 @@ impl Optimizer {
current current
} }
fn is_recursive(
&self,
node: &TypedNode,
global_idx: Option<u32>,
local_slot: Option<u32>,
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( fn visit_node(
&self, &self,
node: TypedNode, node: TypedNode,
@@ -136,16 +173,18 @@ impl Optimizer {
Address::Global(idx) => { Address::Global(idx) => {
if !sub.assigned_globals.contains(&idx) { if !sub.assigned_globals.contains(&idx) {
if let Some(val) = sub.globals.get(&idx) { if let Some(val) = sub.globals.get(&idx) {
return Node { if self.is_inlinable_value(val, Some(idx)) {
identity: node.identity, return Node {
ty: val.static_type(), identity: node.identity,
kind: BoundKind::Constant(val.clone()), ty: val.static_type(),
}; kind: BoundKind::Constant(val.clone()),
};
}
} }
if let Some(globals_rc) = &self.globals { if let Some(globals_rc) = &self.globals {
let globals = globals_rc.borrow(); let globals = globals_rc.borrow();
if let Some(val) = globals.get(idx as usize) if let Some(val) = globals.get(idx as usize)
&& self.is_inlinable_value(val) && self.is_inlinable_value(val, Some(idx))
{ {
return Node { return Node {
identity: node.identity, identity: node.identity,
@@ -225,7 +264,10 @@ impl Optimizer {
&& positional_count.is_some() && positional_count.is_some()
&& path.inlining_depth < 5 && 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; path.inlining_depth += 1;
let collapsed = self.try_beta_reduce_with_sub( let collapsed = self.try_beta_reduce_with_sub(
params, params,
@@ -263,20 +305,8 @@ impl Optimizer {
&& upvalues.is_empty() && upvalues.is_empty()
&& positional_count.is_some() && positional_count.is_some()
{ {
// RECURSION CHECK: Don't inline if the function calls itself // STRICT RECURSION CHECK: Don't inline if recursive (idx or identity)
let mut used_in_body = HashSet::new(); if !self.is_recursive(body, Some(*idx), None, Some(&lambda_node.identity)) {
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) {
let mut inner_sub = SubstitutionMap::new(); let mut inner_sub = SubstitutionMap::new();
path.inlining_stack.insert(*idx); path.inlining_stack.insert(*idx);
path.inlining_depth += 1; path.inlining_depth += 1;
@@ -304,29 +334,39 @@ impl Optimizer {
&& (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(); // STRICT RECURSION CHECK: Don't inline if closure body contains its own identity
for (i, cell) in closure.upvalues.iter().enumerate() { if !self.is_recursive(
closure_sub.add_upvalue(i as u32, cell.borrow().clone()); &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; path.inlining_depth += 1;
let inlined_body = self.visit_node( let inlined_body = self.visit_node(
(*closure.function_node).clone(), (*closure.function_node).clone(),
&mut closure_sub, &mut closure_sub,
path, path,
); );
let collapsed = self.try_beta_reduce_with_sub( 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, path,
); );
path.inlining_depth -= 1; path.inlining_depth -= 1;
path.exit_lambda(&closure.function_node.identity);
if let Some(res) = collapsed { if let Some(res) = collapsed {
return res; return res;
}
} }
} }
@@ -342,6 +382,7 @@ impl Optimizer {
&& (closure.upvalues.is_empty() && (closure.upvalues.is_empty()
|| self.purity_of(&closure.function_node) >= Purity::SideEffectFree) || self.purity_of(&closure.function_node) >= Purity::SideEffectFree)
{ {
// CRACKING fallback (only if not recursive already handled above)
let mut closure_sub = SubstitutionMap::new(); let mut closure_sub = SubstitutionMap::new();
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());
@@ -415,20 +456,11 @@ impl Optimizer {
} }
BoundKind::Block { exprs } => { BoundKind::Block { exprs } => {
let mut used_in_block = HashSet::new(); let mut info = UsageInfo::default();
let mut used_globals_in_block = HashSet::new();
let mut assigned_locals_in_block = HashSet::new();
let mut assigned_globals_in_block = HashSet::new();
if !exprs.is_empty() { if !exprs.is_empty() {
for e in &exprs { for e in &exprs {
self.collect_usage( self.collect_usage(e, &mut info);
e,
&mut used_in_block,
&mut used_globals_in_block,
&mut assigned_locals_in_block,
&mut assigned_globals_in_block,
);
} }
if let Some(last) = exprs.last() { if let Some(last) = exprs.last() {
match &last.kind { match &last.kind {
@@ -436,13 +468,13 @@ impl Optimizer {
addr: Address::Local(slot), addr: Address::Local(slot),
.. ..
} => { } => {
used_in_block.insert(*slot); info.used_locals.insert(*slot);
} }
BoundKind::Get { BoundKind::Get {
addr: Address::Global(idx), 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 // IMPORTANT: Propagate block assignments to the outer map so we don't
// incorrectly inline these variables in the rest of the script. // 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); sub.assigned_locals.insert(*slot);
} }
for idx in &assigned_globals_in_block { for idx in &info.assigned_globals {
sub.assigned_globals.insert(*idx); sub.assigned_globals.insert(*idx);
} }
@@ -467,7 +499,7 @@ impl Optimizer {
if self.enabled && !is_last { if self.enabled && !is_last {
let removable = match &e.kind { let removable = match &e.kind {
BoundKind::DefLocal { slot, value, .. } => { BoundKind::DefLocal { slot, value, .. } => {
!used_in_block.contains(slot) !info.used_locals.contains(slot)
&& !sub.captured_slots.contains(slot) && !sub.captured_slots.contains(slot)
&& self.purity_of(value) >= Purity::SideEffectFree && self.purity_of(value) >= Purity::SideEffectFree
} }
@@ -476,7 +508,7 @@ impl Optimizer {
value, value,
.. ..
} => { } => {
!used_globals_in_block.contains(global_index) !info.used_globals.contains(global_index)
&& (self.purity_of(value) >= Purity::SideEffectFree && (self.purity_of(value) >= Purity::SideEffectFree
|| matches!(value.kind, BoundKind::Lambda { .. })) || matches!(value.kind, BoundKind::Lambda { .. }))
} }
@@ -485,7 +517,7 @@ impl Optimizer {
value, value,
.. ..
} => { } => {
!used_in_block.contains(slot) !info.used_locals.contains(slot)
&& !sub.captured_slots.contains(slot) && !sub.captured_slots.contains(slot)
&& self.purity_of(value) >= Purity::SideEffectFree && self.purity_of(value) >= Purity::SideEffectFree
} }
@@ -538,29 +570,19 @@ impl Optimizer {
unreachable!() unreachable!()
}; };
let mut used_in_lambda = HashSet::new(); let mut info = UsageInfo::default();
let mut used_globals_in_lambda = HashSet::new(); self.collect_usage(&node, &mut info);
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,
);
// Track captures from outer scope that are assigned in this lambda // 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); sub.assigned_globals.insert(*idx);
} }
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.assigned_locals = assigned_locals_in_lambda; next_inner_subs.assigned_locals = info.assigned_locals;
next_inner_subs.assigned_globals = assigned_globals_in_lambda; next_inner_subs.assigned_globals = info.assigned_globals;
for (old_idx, capture_addr) in original_upvalues.iter().enumerate() { for (old_idx, capture_addr) in original_upvalues.iter().enumerate() {
let mut inlined_val = None; let mut inlined_val = None;
@@ -660,7 +682,9 @@ impl Optimizer {
} => { } => {
let value = Box::new(self.visit_node(*value, sub, path)); 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()); if self.is_inlinable_value(val, Some(global_index)) {
sub.add_global(global_index, val.clone());
}
} }
let p = self.purity_of(&value); let p = self.purity_of(&value);
if p > Purity::Impure 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<u32>) -> bool {
match val { match val {
Value::Int(_) Value::Int(_)
| Value::Float(_) | Value::Float(_)
@@ -748,7 +772,14 @@ impl Optimizer {
| Value::DateTime(_) => true, | Value::DateTime(_) => true,
Value::Object(obj) => { Value::Object(obj) => {
if let Some(closure) = obj.as_any().downcast_ref::<Closure>() { if let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
// A closure is only inlinable if it's not recursive
closure.upvalues.is_empty() closure.upvalues.is_empty()
&& !self.is_recursive(
&closure.function_node,
global_idx,
None,
Some(&closure.function_node.identity),
)
} else { } else {
false false
} }
@@ -951,96 +982,46 @@ impl Optimizer {
} }
} }
fn collect_usage( fn collect_usage(&self, node: &TypedNode, info: &mut UsageInfo) {
&self,
node: &TypedNode,
used_locals: &mut HashSet<u32>,
used_globals: &mut HashSet<u32>,
assigned_locals: &mut HashSet<u32>,
assigned_globals: &mut HashSet<u32>,
) {
match &node.kind { match &node.kind {
BoundKind::Constant(v) => { BoundKind::Constant(v) => {
if let Value::Object(obj) = v if let Value::Object(obj) = v
&& let Some(closure) = obj.as_any().downcast_ref::<Closure>() && let Some(closure) = obj.as_any().downcast_ref::<Closure>()
{ {
self.collect_usage( info.used_identities
&closure.function_node, .insert(closure.function_node.identity.clone());
used_locals, self.collect_usage(&closure.function_node, info);
used_globals,
assigned_locals,
assigned_globals,
);
} }
} }
BoundKind::Get { addr, .. } => match addr { BoundKind::Get { addr, .. } => match addr {
Address::Local(slot) => { Address::Local(slot) => {
used_locals.insert(*slot); info.used_locals.insert(*slot);
} }
Address::Global(idx) => { Address::Global(idx) => {
used_globals.insert(*idx); info.used_globals.insert(*idx);
} }
_ => {} _ => {}
}, },
BoundKind::Set { addr, value } => { BoundKind::Set { addr, value } => {
match addr { match addr {
Address::Local(slot) => { Address::Local(slot) => {
assigned_locals.insert(*slot); info.assigned_locals.insert(*slot);
} }
Address::Global(idx) => { Address::Global(idx) => {
assigned_globals.insert(*idx); info.assigned_globals.insert(*idx);
} }
_ => {} _ => {}
} }
self.collect_usage( self.collect_usage(value, info);
value,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
} }
BoundKind::Lambda { BoundKind::Lambda { params, body, .. } => {
params, info.used_identities.insert(node.identity.clone());
body, self.collect_usage(params, info);
upvalues, self.collect_usage(body, info);
..
} => {
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::Block { exprs } => { BoundKind::Block { exprs } => {
for e in exprs { for e in exprs {
self.collect_usage( self.collect_usage(e, info);
e,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
} }
} }
BoundKind::If { BoundKind::If {
@@ -1048,92 +1029,32 @@ impl Optimizer {
then_br, then_br,
else_br, else_br,
} => { } => {
self.collect_usage( self.collect_usage(cond, info);
cond, self.collect_usage(then_br, info);
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
self.collect_usage(
then_br,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
if let Some(e) = else_br { if let Some(e) = else_br {
self.collect_usage( self.collect_usage(e, info);
e,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
} }
} }
BoundKind::Call { callee, args } => { BoundKind::Call { callee, args } => {
self.collect_usage( self.collect_usage(callee, info);
callee, self.collect_usage(args, info);
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
self.collect_usage(
args,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
} }
BoundKind::Tuple { elements } => { BoundKind::Tuple { elements } => {
for e in elements { for e in elements {
self.collect_usage( self.collect_usage(e, info);
e,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
} }
} }
BoundKind::Record { fields } => { BoundKind::Record { fields } => {
for (k, v) in fields { for (k, v) in fields {
self.collect_usage( self.collect_usage(k, info);
k, self.collect_usage(v, info);
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
self.collect_usage(
v,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
} }
} }
BoundKind::DefLocal { value, .. } | BoundKind::DefGlobal { value, .. } => { BoundKind::DefLocal { value, .. } | BoundKind::DefGlobal { value, .. } => {
self.collect_usage( self.collect_usage(value, info);
value,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
} }
BoundKind::Expansion { bound_expanded, .. } => { BoundKind::Expansion { bound_expanded, .. } => {
self.collect_usage( self.collect_usage(bound_expanded, info);
bound_expanded,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
} }
_ => {} _ => {}
} }
src/ast/compiler/snapshots/myc__ast__compiler__optimizer__tests__opt_lambda_cracking.snap
+4 -3
View File
@@ -1,12 +1,13 @@
--- ---
source: src/ast/compiler/optimizer.rs source: src/ast/compiler/optimizer.rs
assertion_line: 727 assertion_line: 1333
expression: dump expression: dump
--- ---
Lambda (Upvalues: 0) <Metadata: Metadata { ty: Function(Signature { params: Tuple([]), ret: Function(Signature { params: Tuple([]), ret: Int }) }), is_tail: true }> Lambda (Upvalues: 0) <Metadata: Metadata { ty: Function(Signature { params: Tuple([]), ret: Function(Signature { params: Tuple([]), ret: Int }) }), is_tail: true }>
Parameters: Parameters:
Tuple <Metadata: Metadata { ty: Tuple([]), is_tail: false }> Tuple <Metadata: Metadata { ty: Tuple([]), is_tail: false }>
Lambda (Upvalues: 0) <Metadata: Metadata { ty: Function(Signature { params: Tuple([]), ret: Int }), is_tail: true }> Lambda (Upvalues: 1) <Metadata: Metadata { ty: Function(Signature { params: Tuple([]), ret: Int }), is_tail: true }>
Parameters: Parameters:
Tuple <Metadata: Metadata { ty: Tuple([]), is_tail: false }> Tuple <Metadata: Metadata { ty: Tuple([]), is_tail: false }>
Constant: 10 <Metadata: Metadata { ty: Int, is_tail: true }> Upvalues: [Local(0)]
Get: x (Upvalue(0)) <Metadata: Metadata { ty: Int, is_tail: true }>