From 283cdf61a4e45e0aa86334d48ef297df51c28999 Mon Sep 17 00:00:00 2001 From: Michael Schimmel Date: Wed, 25 Feb 2026 10:57:58 +0100 Subject: [PATCH] Refactor UsageInfo and SubstitutionMap This commit refactors the `UsageInfo` struct to use a single `HashSet
` for both used and assigned addresses, simplifying the logic. It also updates the `SubstitutionMap` to use a similar approach for tracking assigned values. Key changes include: - `UsageInfo` now has `used` and `assigned` fields of type `HashSet
`. - `SubstitutionMap`'s `add_local`, `add_global`, `add_upvalue`, `remove_local`, `remove_global`, and `remove_upvalue` methods have been replaced with a generic `add_value` and `remove_value` that operate on `Address`. - The `Optimizer`'s `collect_pattern_usage` and `visit_node` methods have been updated to use the new `UsageInfo` and `SubstitutionMap` APIs. - `Get` and `Set` nodes now use `sub.map_address` to transform addresses based on the current substitution. --- src/ast/compiler/binder.rs | 15 +- src/ast/compiler/bound_nodes.rs | 4 +- src/ast/compiler/optimizer.rs | 506 +++++++++++++------------------ src/ast/compiler/type_checker.rs | 8 +- 4 files changed, 234 insertions(+), 299 deletions(-) diff --git a/src/ast/compiler/binder.rs b/src/ast/compiler/binder.rs index 9e4dc37..f1a8d3c 100644 --- a/src/ast/compiler/binder.rs +++ b/src/ast/compiler/binder.rs @@ -118,7 +118,10 @@ impl Binder { if self.functions.len() == 1 { let mut globals = self.globals.borrow_mut(); if globals.contains_key(name) { - return Err(format!("Global variable '{}' is already defined.", name.name)); + return Err(format!( + "Global variable '{}' is already defined.", + name.name + )); } let idx = globals.len() as u32; globals.insert(name.clone(), idx); @@ -512,7 +515,10 @@ mod tests { if let BoundKind::Lambda { body, .. } = &bound.kind { if let BoundKind::Block { exprs } = &body.kind { let x_decl = &exprs[0]; - if let BoundKind::Define { captured_by, addr, .. } = &x_decl.kind { + if let BoundKind::Define { + captured_by, addr, .. + } = &x_decl.kind + { assert!(matches!(addr, Address::Local(_))); assert!( !captured_by.is_empty(), @@ -544,7 +550,10 @@ mod tests { if let BoundKind::Lambda { body, .. } = &bound.kind { if let BoundKind::Block { exprs } = &body.kind { let x_decl = &exprs[0]; - if let BoundKind::Define { captured_by, addr, .. } = &x_decl.kind { + if let BoundKind::Define { + captured_by, addr, .. + } = &x_decl.kind + { assert!(matches!(addr, Address::Local(_))); assert!( captured_by.is_empty(), diff --git a/src/ast/compiler/bound_nodes.rs b/src/ast/compiler/bound_nodes.rs index ee00e00..61bc0d2 100644 --- a/src/ast/compiler/bound_nodes.rs +++ b/src/ast/compiler/bound_nodes.rs @@ -237,7 +237,9 @@ impl BoundKind { BoundKind::Constant(v) => format!("CONST({})", v), BoundKind::Get { addr, name } => format!("GET({}, {:?})", name.name, addr), BoundKind::Set { addr, .. } => format!("SET({:?})", addr), - BoundKind::Define { name, addr, kind, .. } => { + BoundKind::Define { + name, addr, kind, .. + } => { let k_str = match kind { DeclarationKind::Variable => "VAR", DeclarationKind::Parameter => "PARAM", diff --git a/src/ast/compiler/optimizer.rs b/src/ast/compiler/optimizer.rs index 449a269..e13ae2c 100644 --- a/src/ast/compiler/optimizer.rs +++ b/src/ast/compiler/optimizer.rs @@ -44,13 +44,19 @@ impl PathTracker { #[derive(Default)] struct UsageInfo { - used_locals: HashSet, - used_globals: HashSet, - used_upvalues: HashSet, + used: HashSet
, + assigned: HashSet
, used_identities: HashSet, - assigned_locals: HashSet, - assigned_globals: HashSet, - assigned_upvalues: HashSet, +} + +impl UsageInfo { + fn is_assigned(&self, addr: &Address) -> bool { + self.assigned.contains(addr) + } + + fn is_used(&self, addr: &Address) -> bool { + self.used.contains(addr) + } } impl Optimizer { @@ -99,26 +105,12 @@ impl Optimizer { fn collect_pattern_usage(&self, node: &AnalyzedNode, info: &mut UsageInfo) { match &node.kind { - BoundKind::Define { addr, .. } => match addr { - Address::Local(slot) => { - info.assigned_locals.insert(*slot); - } - Address::Global(idx) => { - info.assigned_globals.insert(*idx); - } - _ => {} - }, - BoundKind::Set { addr, .. } => 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::Define { addr, .. } => { + info.assigned.insert(*addr); + } + BoundKind::Set { addr, .. } => { + info.assigned.insert(*addr); + } BoundKind::Tuple { elements } => { for el in elements { self.collect_pattern_usage(el, info); @@ -136,54 +128,65 @@ impl Optimizer { ) -> AnalyzedNode { let (new_kind, metrics) = match node.kind { BoundKind::Get { addr, ref name } => { - let new_addr = match addr { - Address::Local(slot) => { - if !sub.assigned_locals.contains(&slot) { - if let Some(inlined_node) = sub.ast_locals.get(&slot) { - return inlined_node.clone(); - } - if let Some(val) = sub.locals.get(&slot) { - return self.make_constant_node(val.clone(), &node); - } - } - sub.used_slots.insert(slot); - Address::Local(sub.map_slot(slot)) - } - Address::Upvalue(idx) => { - if let Some(val) = sub.upvalues.get(&idx) { + if !sub.assigned.contains(&addr) { + // 1. Try inlining from current substitution map (locals/globals/upvalues) + if let Some(val) = sub.get_value(&addr) { + // For globals, only inline if it's a known "safe" value + let can_inline = match addr { + Address::Global(idx) => self.is_inlinable_value(val, Some(idx)), + _ => true, + }; + if can_inline { return self.make_constant_node(val.clone(), &node); } - Address::Upvalue(idx) } - Address::Global(idx) => { - if !sub.assigned_globals.contains(&idx) { - if let Some(val) = sub.globals.get(&idx) - && self.is_inlinable_value(val, Some(idx)) - { - return self.make_constant_node(val.clone(), &node); - } - 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, Some(idx)) - { - return self.make_constant_node(val.clone(), &node); - } - } + + // 2. Special case: Local-only AST inlining (for non-constant but pure expressions) + if let Address::Local(slot) = addr + && let Some(inlined_node) = sub.ast_locals.get(&slot) + { + return inlined_node.clone(); + } + + // 3. Fallback for Globals: check the actual VM environment + if let Address::Global(idx) = addr + && 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, Some(idx)) + { + return self.make_constant_node(val.clone(), &node); } - sub.used_globals.insert(idx); - Address::Global(idx) } - }; + } + + sub.used.insert(addr); ( BoundKind::Get { - addr: new_addr, + addr: sub.map_address(addr), name: name.clone(), }, node.ty.clone(), ) } + BoundKind::Set { addr, value } => { + let value = Box::new(self.visit_node(*value, sub, path)); + if let BoundKind::Constant(val) = &value.kind { + sub.add_value(addr, val.clone()); + } else { + sub.remove_value(&addr); + } + ( + BoundKind::Set { + addr: sub.map_address(addr), + value, + }, + node.ty.clone(), + ) + } + BoundKind::Define { ref name, addr, @@ -191,68 +194,41 @@ impl Optimizer { ref value, ref captured_by, } => { - let value = Box::new(self.visit_node(*value.clone(), sub, path)); - let new_addr = match addr { - Address::Local(slot) => { - if let BoundKind::Constant(val) = &value.kind { - sub.add_local(slot, val.clone()); - } else { - sub.locals.remove(&slot); - } - Address::Local(sub.map_slot(slot)) - } - Address::Global(idx) => { - if let BoundKind::Constant(val) = &value.kind { - sub.add_global(idx, val.clone()); - } else { - sub.globals.remove(&idx); - } - Address::Global(idx) - } - _ => addr, - }; + let value_opt = Box::new(self.visit_node((**value).clone(), sub, path)); + + if let Address::Local(slot) = addr + && !captured_by.is_empty() + { + sub.captured_slots.insert(slot); + } + + if let BoundKind::Constant(val) = &value_opt.kind { + sub.add_value(addr, val.clone()); + } else { + sub.remove_value(&addr); + } + + if let Address::Global(global_index) = addr + && value_opt.ty.purity > Purity::Impure + && let Some(purity_rc) = &self.global_purity + { + purity_rc + .borrow_mut() + .insert(global_index, value_opt.ty.purity); + } + ( BoundKind::Define { name: name.clone(), - addr: new_addr, + addr: sub.map_address(addr), kind, - value, + value: value_opt, captured_by: captured_by.clone(), }, node.ty.clone(), ) } - BoundKind::Set { addr, value } => { - 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 { - sub.add_local(slot, val.clone()); - } else { - sub.locals.remove(&slot); - } - Address::Local(sub.map_slot(slot)) - } - Address::Global(idx) => { - if let BoundKind::Constant(val) = &value.kind { - sub.add_global(idx, val.clone()); - } else { - sub.globals.remove(&idx); - } - Address::Global(idx) - } - _ => addr, - }; - ( - BoundKind::Set { - addr: new_addr, - value, - }, - node.ty.clone(), - ) - } - BoundKind::Call { callee, args } => { let callee = self.visit_node(*callee, sub, path); let args = self.visit_node(*args, sub, path); @@ -335,7 +311,8 @@ impl Optimizer { { let mut closure_sub = SubstitutionMap::new(); for (i, cell) in closure.upvalues.iter().enumerate() { - closure_sub.add_upvalue(i as u32, cell.borrow().clone()); + closure_sub + .add_value(Address::Upvalue(i as u32), cell.borrow().clone()); } path.inlining_depth += 1; @@ -424,12 +401,7 @@ impl Optimizer { } } - for slot in &info.assigned_locals { - sub.assigned_locals.insert(*slot); - } - for idx in &info.assigned_globals { - sub.assigned_globals.insert(*idx); - } + sub.assigned.extend(info.assigned.iter().cloned()); let mut new_exprs = Vec::with_capacity(exprs.len()); let last_idx = exprs.len().saturating_sub(1); @@ -438,26 +410,23 @@ impl Optimizer { let is_last = i == last_idx; if self.enabled && !is_last { let removable = match &e.kind { - BoundKind::Define { addr, value, .. } => match addr { - Address::Local(slot) => { - !info.used_locals.contains(slot) - && !sub.captured_slots.contains(slot) - && value.ty.purity >= Purity::SideEffectFree - } - Address::Global(global_index) => { - !info.used_globals.contains(global_index) - && (value.ty.purity >= Purity::SideEffectFree - || matches!(value.kind, BoundKind::Lambda { .. })) - } - _ => false, - }, - BoundKind::Set { - addr: Address::Local(slot), - value, - .. - } => { - !info.used_locals.contains(slot) - && !sub.captured_slots.contains(slot) + BoundKind::Define { addr, value, .. } => { + !info.is_used(addr) + && (if let Address::Local(slot) = addr { + !sub.captured_slots.contains(slot) + } else { + true + }) + && (value.ty.purity >= Purity::SideEffectFree + || matches!(value.kind, BoundKind::Lambda { .. })) + } + BoundKind::Set { addr, value, .. } => { + !info.is_used(addr) + && (if let Address::Local(slot) = addr { + !sub.captured_slots.contains(slot) + } else { + true + }) && value.ty.purity >= Purity::SideEffectFree } _ => e.ty.purity >= Purity::SideEffectFree, @@ -503,44 +472,26 @@ impl Optimizer { let mut info = UsageInfo::default(); self.collect_usage(&node, &mut info); - 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 = info.assigned_locals; - next_inner_subs.assigned_globals = info.assigned_globals; + next_inner_subs.assigned = info.assigned; for (old_idx, capture_addr) in original_upvalues.iter().enumerate() { let mut inlined_val = None; - match capture_addr { - Address::Local(slot) => { - if let Some(val) = sub.locals.get(slot) { - inlined_val = Some(val.clone()); - } - sub.captured_slots.insert(*slot); - } - Address::Upvalue(idx) => { - if let Some(val) = sub.upvalues.get(idx) { - inlined_val = Some(val.clone()); - } - } - _ => {} + if let Some(val) = sub.get_value(capture_addr) { + inlined_val = Some(val.clone()); + } + if let Address::Local(slot) = capture_addr { + sub.captured_slots.insert(*slot); } if let Some(val) = inlined_val { - next_inner_subs.add_upvalue(old_idx as u32, val); + next_inner_subs.add_value(Address::Upvalue(old_idx as u32), val); mapping.push(None); } else { mapping.push(Some(new_upvalues.len() as u32)); - let new_addr = if let Address::Local(parent_slot) = capture_addr { - Address::Local(sub.map_slot(*parent_slot)) - } else { - *capture_addr - }; - new_upvalues.push(new_addr); + new_upvalues.push(sub.map_address(*capture_addr)); } } @@ -575,11 +526,8 @@ impl Optimizer { let mut info = UsageInfo::default(); self.collect_pattern_usage(&pat_opt, &mut info); - for slot in info.assigned_locals { - sub.locals.remove(&slot); - } - for idx in info.assigned_globals { - sub.globals.remove(&idx); + for addr in info.assigned { + sub.remove_value(&addr); } ( @@ -633,7 +581,10 @@ impl Optimizer { fn collect_parameter_slots(&self, node: &AnalyzedNode, sub: &mut SubstitutionMap) { match &node.kind { - BoundKind::Define { addr: Address::Local(slot), .. } => { + BoundKind::Define { + addr: Address::Local(slot), + .. + } => { sub.slot_mapping.insert(*slot, *slot); sub.next_slot = sub.next_slot.max(*slot + 1); } @@ -752,26 +703,20 @@ impl Optimizer { } // 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) + let addr = Address::Local(slot); + if (body_usage.is_used(&addr) || body_usage.is_assigned(&addr)) + && sub.get_value(&addr).is_none() && !sub.ast_locals.contains_key(&slot) { return None; } } - if sub.locals.is_empty() - && sub.ast_locals.is_empty() - && sub.upvalues.is_empty() - && !arg_vals.is_empty() - { + if sub.values.is_empty() && sub.ast_locals.is_empty() && !arg_vals.is_empty() { return None; } @@ -780,7 +725,10 @@ impl Optimizer { fn collect_parameter_slots_set(&self, node: &AnalyzedNode, slots: &mut HashSet) { match &node.kind { - BoundKind::Define { addr: Address::Local(slot), .. } => { + BoundKind::Define { + addr: Address::Local(slot), + .. + } => { slots.insert(*slot); } BoundKind::Tuple { elements } => { @@ -818,26 +766,32 @@ impl Optimizer { body_usage: &UsageInfo, ) { match &pattern.kind { - BoundKind::Define { addr: Address::Local(slot), .. } => { + BoundKind::Define { addr, .. } => { if let Some(arg) = args.get(*offset) - && !body_usage.assigned_locals.contains(slot) + && !body_usage.is_assigned(addr) { 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()); - } + sub.add_value(*addr, val.clone()); + } else if let BoundKind::Lambda { upvalues, .. } = &arg.kind + && upvalues.is_empty() + && let Address::Local(slot) = addr + { + sub.add_ast_local(*slot, arg.clone()); } else if let BoundKind::Get { addr: Address::Global(_), .. } = &arg.kind + && let Address::Local(slot) = addr { sub.add_ast_local(*slot, arg.clone()); } } - sub.slot_mapping.insert(*slot, *slot); - sub.next_slot = sub.next_slot.max(*slot + 1); + + // Ensure local slots are mapped even if no constant/AST is substituted + if let Address::Local(slot) = addr { + sub.slot_mapping.insert(*slot, *slot); + sub.next_slot = sub.next_slot.max(*slot + 1); + } *offset += 1; } BoundKind::Tuple { elements } => { @@ -894,29 +848,11 @@ impl Optimizer { self.collect_usage(&closure.function_node, info); } } - BoundKind::Get { addr, .. } => 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); - } - }, + BoundKind::Get { addr, .. } => { + info.used.insert(*addr); + } BoundKind::Set { addr, value } => { - 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); - } - } + info.assigned.insert(*addr); self.collect_usage(value, info); } BoundKind::Lambda { @@ -932,40 +868,32 @@ impl Optimizer { 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); + for addr in &inner_info.used { + if let Address::Global(_) = addr { + info.used.insert(*addr); + } + } + for addr in &inner_info.assigned { + if let Address::Global(_) = addr { + info.assigned.insert(*addr); + } + } 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); - } - } + for addr in &inner_info.used { + if let Address::Upvalue(idx) = addr + && let Some(parent_addr) = upvalues.get(*idx as usize) + { + info.used.insert(*parent_addr); } } // 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); - } - } + for addr in &inner_info.assigned { + if let Address::Upvalue(idx) = addr + && let Some(parent_addr) = upvalues.get(*idx as usize) + { + info.assigned.insert(*parent_addr); } } } @@ -1012,32 +940,24 @@ impl Optimizer { } struct SubstitutionMap { - locals: HashMap, - upvalues: HashMap, - globals: HashMap, + values: HashMap, ast_locals: HashMap, slot_mapping: HashMap, - assigned_locals: HashSet, - assigned_globals: HashSet, + assigned: HashSet
, next_slot: u32, - used_slots: HashSet, - used_globals: HashSet, + used: HashSet
, captured_slots: HashSet, } impl SubstitutionMap { fn new() -> Self { Self { - locals: HashMap::new(), - upvalues: HashMap::new(), - globals: HashMap::new(), + values: HashMap::new(), ast_locals: HashMap::new(), slot_mapping: HashMap::new(), - assigned_locals: HashSet::new(), - assigned_globals: HashSet::new(), + assigned: HashSet::new(), next_slot: 0, - used_slots: HashSet::new(), - used_globals: HashSet::new(), + used: HashSet::new(), captured_slots: HashSet::new(), } } @@ -1045,6 +965,7 @@ impl SubstitutionMap { fn add_ast_local(&mut self, slot: u32, node: AnalyzedNode) { self.ast_locals.insert(slot, node); } + fn map_slot(&mut self, old_slot: u32) -> u32 { if let Some(&new_slot) = self.slot_mapping.get(&old_slot) { return new_slot; @@ -1054,42 +975,46 @@ impl SubstitutionMap { self.next_slot += 1; new_slot } - fn add_local(&mut self, slot: u32, val: Value) { - self.locals.insert(slot, val); + + fn map_address(&mut self, addr: Address) -> Address { + match addr { + Address::Local(slot) => Address::Local(self.map_slot(slot)), + other => other, + } } - fn add_upvalue(&mut self, idx: u32, val: Value) { - self.upvalues.insert(idx, val); + + fn add_value(&mut self, addr: Address, val: Value) { + self.values.insert(addr, val); } - fn add_global(&mut self, idx: u32, val: Value) { - self.globals.insert(idx, val); + + fn get_value(&self, addr: &Address) -> Option<&Value> { + self.values.get(addr) + } + + fn remove_value(&mut self, addr: &Address) { + self.values.remove(addr); + } + + fn reindex_addr(&self, addr: Address, mapping: &[Option]) -> Address { + if let Address::Upvalue(idx) = addr + && let Some(res) = mapping.get(idx as usize) + && let Some(new_idx) = res + { + Address::Upvalue(*new_idx) + } else { + addr + } } fn reindex_upvalues(&self, node: AnalyzedNode, mapping: &[Option]) -> AnalyzedNode { let (new_kind, metrics) = match node.kind { - BoundKind::Get { - addr: Address::Upvalue(idx), - name, - } => { - if let Some(res) = mapping.get(idx as usize) - && let Some(new_idx) = res - { - ( - BoundKind::Get { - addr: Address::Upvalue(*new_idx), - name, - }, - node.ty.clone(), - ) - } else { - ( - BoundKind::Get { - addr: Address::Upvalue(idx), - name, - }, - node.ty.clone(), - ) - } - } + BoundKind::Get { addr, name } => ( + BoundKind::Get { + addr: self.reindex_addr(addr, mapping), + name, + }, + node.ty.clone(), + ), BoundKind::Lambda { params, upvalues, @@ -1098,14 +1023,7 @@ impl SubstitutionMap { } => { let mut next_upvalues = Vec::new(); for addr in upvalues { - if let Address::Upvalue(idx) = addr - && let Some(res) = mapping.get(idx as usize) - && let Some(new_idx) = res - { - next_upvalues.push(Address::Upvalue(*new_idx)); - continue; - } - next_upvalues.push(addr); + next_upvalues.push(self.reindex_addr(addr, mapping)); } ( BoundKind::Lambda { diff --git a/src/ast/compiler/type_checker.rs b/src/ast/compiler/type_checker.rs index aedc980..31d8adb 100644 --- a/src/ast/compiler/type_checker.rs +++ b/src/ast/compiler/type_checker.rs @@ -294,7 +294,13 @@ impl TypeChecker { } else { ctx.get_type(addr) }; - (BoundKind::Get { addr, name: name.clone() }, ty) + ( + BoundKind::Get { + addr, + name: name.clone(), + }, + ty, + ) } BoundKind::Set { addr, value } => {