Refactor optimizer to track upvalue usage
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
+108
-19
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user