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.
This commit is contained in:
Michael Schimmel
+131
-210
@@ -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 {
|
||||
pub fn new(enabled: bool) -> Self {
|
||||
Self {
|
||||
@@ -98,6 +107,34 @@ impl Optimizer {
|
||||
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(
|
||||
&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<u32>) -> 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::<Closure>() {
|
||||
// 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<u32>,
|
||||
used_globals: &mut HashSet<u32>,
|
||||
assigned_locals: &mut HashSet<u32>,
|
||||
assigned_globals: &mut HashSet<u32>,
|
||||
) {
|
||||
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::<Closure>()
|
||||
{
|
||||
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);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
+4
-3
@@ -1,12 +1,13 @@
|
||||
---
|
||||
source: src/ast/compiler/optimizer.rs
|
||||
assertion_line: 727
|
||||
assertion_line: 1333
|
||||
expression: dump
|
||||
---
|
||||
Lambda (Upvalues: 0) <Metadata: Metadata { ty: Function(Signature { params: Tuple([]), ret: Function(Signature { params: Tuple([]), ret: Int }) }), is_tail: true }>
|
||||
Parameters:
|
||||
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:
|
||||
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 }>
|
||||
|
||||
Reference in New Issue
Block a user