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

Refactor optimizer to track upvalue usage

Browse Source 
The `UsageInfo` struct has been updated to include tracking for
`used_upvalues` and `assigned_upvalues`. The `collect_usage` function
has been modified to handle these new fields when encountering
`Address::Upvalue`.

The `map_params_to_args` function now takes `body_usage` as an argument
to ensure that parameters assigned to or used in the body are correctly
substituted. A new helper function, `collect_parameter_slots_set`, has
been added to gather all parameter slots within a pattern.

The inlining logic in `Optimizer::inline_call` has been enhanced to
prevent inlining if a parameter is used or assigned in the function body
but cannot be substituted. This addresses a bug where aggressive
inlining could lead to incorrect code generation when parameters were
reassigned.

A new integration test, `test_closure_reassignment_optimization_bug`,
has been added to specifically target and verify the fix for this
inlining issue.
This commit is contained in:
Michael Schimmel
2026-02-24 12:53:06 +01:00
parent 2c652e0140
commit e07acc0208
4 changed files with 121 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ast/compiler/optimizer.rs
+108 -19
View File
@@ -46,9 +46,11 @@ impl PathTracker {
struct UsageInfo {
used_locals: HashSet<u32>,
used_globals: HashSet<u32>,
used_upvalues: HashSet<u32>,
used_identities: HashSet<crate::ast::types::Identity>,
assigned_locals: HashSet<u32>,
assigned_globals: HashSet<u32>,
assigned_upvalues: HashSet<u32>,
}
impl Optimizer {
@@ -110,7 +112,9 @@ impl Optimizer {
Address::Global(idx) => {
info.assigned_globals.insert(*idx);
}
_ => {}
Address::Upvalue(idx) => {
info.assigned_upvalues.insert(*idx);
}
},
BoundKind::Tuple { elements } => {
for el in elements {
@@ -759,15 +763,34 @@ impl Optimizer {
sub: &mut SubstitutionMap,
path: &mut PathTracker,
) -> Option<AnalyzedNode> {
let mut body_usage = UsageInfo::default();
self.collect_usage(&body, &mut body_usage);
let mut arg_vals = Vec::new();
self.flatten_tuple(args.clone(), &mut arg_vals);
let mut slot_index = 0;
self.map_params_to_args(params, &arg_vals, &mut slot_index, sub);
self.map_params_to_args(params, &arg_vals, &mut slot_index, sub, &body_usage);
if slot_index != arg_vals.len() {
return None;
}
// Verify that all used/assigned parameters have been substituted.
// If a parameter is used in the body but we couldn't substitute it (e.g. because it's assigned),
// we cannot inline the function because the parameter slot would become invalid.
let mut param_slots = HashSet::new();
self.collect_parameter_slots_set(params, &mut param_slots);
for slot in param_slots {
if (body_usage.used_locals.contains(&slot) || body_usage.assigned_locals.contains(&slot))
&& !sub.locals.contains_key(&slot)
&& !sub.ast_locals.contains_key(&slot)
{
return None;
}
}
if sub.locals.is_empty()
&& sub.ast_locals.is_empty()
&& sub.upvalues.is_empty()
@@ -779,6 +802,20 @@ impl Optimizer {
Some(self.visit_node(body, sub, path))
}
fn collect_parameter_slots_set(&self, node: &AnalyzedNode, slots: &mut HashSet<u32>) {
match &node.kind {
BoundKind::Parameter { slot, .. } => {
slots.insert(*slot);
}
BoundKind::Tuple { elements } => {
for el in elements {
self.collect_parameter_slots_set(el, slots);
}
}
_ => {}
}
}
fn flatten_tuple(&self, node: AnalyzedNode, into: &mut Vec<AnalyzedNode>) {
match node.kind {
BoundKind::Tuple { elements } => {
@@ -802,22 +839,25 @@ impl Optimizer {
args: &[AnalyzedNode],
offset: &mut usize,
sub: &mut SubstitutionMap,
body_usage: &UsageInfo,
) {
match &pattern.kind {
BoundKind::Parameter { slot, .. } => {
if let Some(arg) = args.get(*offset) {
if let BoundKind::Constant(val) = &arg.kind {
sub.add_local(*slot, val.clone());
} else if let BoundKind::Lambda { upvalues, .. } = &arg.kind {
if upvalues.is_empty() {
if !body_usage.assigned_locals.contains(slot) {
if let BoundKind::Constant(val) = &arg.kind {
sub.add_local(*slot, val.clone());
} else if let BoundKind::Lambda { upvalues, .. } = &arg.kind {
if upvalues.is_empty() {
sub.add_ast_local(*slot, arg.clone());
}
} else if let BoundKind::Get {
addr: Address::Global(_),
..
} = &arg.kind
{
sub.add_ast_local(*slot, arg.clone());
}
} else if let BoundKind::Get {
addr: Address::Global(_),
..
} = &arg.kind
{
sub.add_ast_local(*slot, arg.clone());
}
}
sub.slot_mapping.insert(*slot, *slot);
@@ -841,7 +881,7 @@ impl Optimizer {
// Match inner elements against the flattened compound argument
let mut sub_offset = 0;
for el in elements {
self.map_params_to_args(el, &sub_args, &mut sub_offset, sub);
self.map_params_to_args(el, &sub_args, &mut sub_offset, sub, body_usage);
}
*offset += 1;
return;
@@ -850,7 +890,7 @@ impl Optimizer {
// Fallback: Continue flat matching (original behavior)
for el in elements {
self.map_params_to_args(el, args, offset, sub);
self.map_params_to_args(el, args, offset, sub, body_usage);
}
}
_ => {}
@@ -879,7 +919,9 @@ impl Optimizer {
Address::Global(idx) => {
info.used_globals.insert(*idx);
}
_ => {}
Address::Upvalue(idx) => {
info.used_upvalues.insert(*idx);
}
},
BoundKind::Set { addr, value } => {
match addr {
@@ -889,14 +931,61 @@ impl Optimizer {
Address::Global(idx) => {
info.assigned_globals.insert(*idx);
}
_ => {}
Address::Upvalue(idx) => {
info.assigned_upvalues.insert(*idx);
}
}
self.collect_usage(value, info);
}
BoundKind::Lambda { params, body, .. } => {
BoundKind::Lambda {
params,
body,
upvalues,
..
} => {
info.used_identities.insert(node.identity.clone());
self.collect_usage(params, info);
self.collect_usage(body, info);
let mut inner_info = UsageInfo::default();
self.collect_usage(params, &mut inner_info);
self.collect_usage(body, &mut inner_info);
// Propagate globals and identities
info.used_globals.extend(inner_info.used_globals);
info.assigned_globals.extend(inner_info.assigned_globals);
info.used_identities.extend(inner_info.used_identities);
// Map used upvalues to parent scope
for idx in inner_info.used_upvalues {
if let Some(addr) = upvalues.get(idx as usize) {
match addr {
Address::Local(slot) => {
info.used_locals.insert(*slot);
}
Address::Global(idx) => {
info.used_globals.insert(*idx);
}
Address::Upvalue(idx) => {
info.used_upvalues.insert(*idx);
}
}
}
}
// Map assigned upvalues to parent scope
for idx in inner_info.assigned_upvalues {
if let Some(addr) = upvalues.get(idx as usize) {
match addr {
Address::Local(slot) => {
info.assigned_locals.insert(*slot);
}
Address::Global(idx) => {
info.assigned_globals.insert(*idx);
}
Address::Upvalue(idx) => {
info.assigned_upvalues.insert(*idx);
}
}
}
}
}
BoundKind::Block { exprs } => {
for e in exprs {
src/integration_test.rs
+11
View File
@@ -390,4 +390,15 @@ mod tests {
"[\"Symbol:\" \"btc\" \"field:\" :close \"id:\" \"cls\"]"
);
}
#[test]
fn test_closure_reassignment_optimization_bug() {
let env = Environment::new();
// This test case reproduces a bug where the optimizer aggressively inlined a function
// ('f') even though its parameters ('x') were being assigned to in the body (by inner lambda).
// The fix ensures that such functions are NOT inlined.
let source = "(do (def f (fn [[x y]] (fn [] (assign x (+ x y))))) ((f [1 2])))";
let res = env.run_script(source);
assert_eq!(format!("{}", res.unwrap()), "3");
}
}