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

Refactor: Improve optimizer and DCE

Browse Source 
This commit introduces several improvements to the AST optimizer and
dead code elimination (DCE) logic:

- **Address Mapping:** The `Get` and `Set` operations now correctly map
  local slots and global indices using the substitution map, ensuring
  that remapped variables are handled properly.
- **Parameter Slot Mapping:** Parameters are now explicitly mapped to
  their correct slots within the `Lambda` node, preventing potential
  slot reassignments in the function body.
- **Pure Function Analysis:** The `is_pure` function has been refined to
  accurately identify pure expressions, including calls to pure
  functions and stateless closures.
- **Dead Code Elimination:** DCE logic in `BoundKind::Block` has been
  enhanced to remove unused local definitions and pure, non-essential
  expressions more effectively. Global DCE is also improved.
- **Constant Folding:** `try_fold_pure` is now more robust in folding
  pure function calls with constant arguments, reducing redundant
  computations.
- **Beta Reduction:** `try_beta_reduce` has been updated to avoid
  reducing if the body contains `DefLocal`, ensuring correctness.
- **Global Purity Tracking:** A `global_purity` map is introduced to
  track the purity of global values, enabling better optimization of
  global accesses.
- **Inlinable Value Check:** The `is_inlinable_value` check now
  correctly identifies stateless closures for inlining.
- **Optimization Level:** The default `optimization_level` is set to 2
  to enable more aggressive optimizations.
This commit is contained in:
Michael Schimmel
2026-02-22 00:41:02 +01:00
parent 0107264026
commit e87f7232e6
2 changed files with 175 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ast/compiler/optimizer.rs
+174 -154
View File
@@ -50,7 +50,7 @@ impl Optimizer {
fn visit_node(&self, node: TypedNode, sub: &mut SubstitutionMap) -> TypedNode {
let (new_kind, new_ty) = match node.kind {
BoundKind::Get { addr, name } => {
match addr {
let new_addr = match addr {
Address::Local(slot) => {
if let Some(val) = sub.locals.get(&slot) {
return Node {
@@ -60,6 +60,7 @@ impl Optimizer {
};
}
sub.used_slots.insert(slot);
Address::Local(sub.map_slot(slot))
}
Address::Upvalue(idx) => {
if let Some(val) = sub.upvalues.get(&idx) {
@@ -69,10 +70,9 @@ impl Optimizer {
kind: BoundKind::Constant(val.clone()),
};
}
Address::Upvalue(idx)
}
Address::Global(idx) => {
// Phase 4: Global Inlining
// 1. Check script-local global substitution map
if let Some(val) = sub.globals.get(&idx) {
return Node {
identity: node.identity,
@@ -80,11 +80,9 @@ impl Optimizer {
kind: BoundKind::Constant(val.clone()),
};
}
// 2. Check environment global values
if let Some(globals_rc) = &self.globals {
let globals = globals_rc.borrow();
if let Some(val) = globals.get(idx as usize) {
// Only inline "stable" values: scalars and stateless functions
if self.is_inlinable_value(val) {
return Node {
identity: node.identity,
@@ -95,11 +93,41 @@ impl Optimizer {
}
}
sub.used_globals.insert(idx);
Address::Global(idx)
}
}
(BoundKind::Get { addr, name }, node.ty)
};
(BoundKind::Get { addr: new_addr, name }, node.ty)
}
BoundKind::Parameter { name, slot } => {
let new_slot = sub.map_slot(slot);
(BoundKind::Parameter { name, slot: new_slot }, node.ty)
}
BoundKind::Set { addr, value } => {
let value = Box::new(self.visit_node(*value, sub));
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)
},
BoundKind::Call { callee, args } => {
let callee = self.visit_node(*callee, sub);
let args = self.visit_node(*args, sub);
@@ -178,44 +206,48 @@ impl Optimizer {
},
BoundKind::Block { exprs } => {
let mut used_in_block = HashSet::new();
if !exprs.is_empty() {
for e in &exprs {
self.collect_local_usage(e, &mut used_in_block);
}
if let Some(last) = exprs.last() {
if let BoundKind::Get { addr: Address::Local(slot), .. } = &last.kind {
used_in_block.insert(*slot);
}
}
}
let mut new_exprs = Vec::with_capacity(exprs.len());
for e in exprs {
let last_idx = exprs.len().saturating_sub(1);
for (i, e) in exprs.into_iter().enumerate() {
let is_last = i == last_idx;
if self.level >= 2 && !is_last {
let removable = match &e.kind {
BoundKind::DefLocal { slot, value, .. } => {
!used_in_block.contains(slot) && !sub.captured_slots.contains(slot) && self.is_pure(value)
}
BoundKind::Set { addr: Address::Local(slot), value, .. } => {
!used_in_block.contains(slot) && !sub.captured_slots.contains(slot) && self.is_pure(value)
}
_ => self.is_pure(&e),
};
if removable {
continue;
}
}
let opt = self.visit_node(e, sub);
if self.level >= 2 && matches!(opt.kind, BoundKind::Nop) {
if self.level >= 2 && matches!(opt.kind, BoundKind::Nop) && !is_last {
continue;
}
new_exprs.push(opt);
}
if self.level >= 2 {
if !new_exprs.is_empty() {
let last_idx = new_exprs.len() - 1;
let mut filtered = Vec::with_capacity(new_exprs.len());
for (i, e) in new_exprs.into_iter().enumerate() {
if i == last_idx {
filtered.push(e);
continue;
}
let removable = match &e.kind {
BoundKind::DefLocal { slot, .. } | BoundKind::Set { addr: Address::Local(slot), .. } => {
!sub.used_slots.contains(slot) && !sub.captured_slots.contains(slot) && self.is_pure(&e)
}
BoundKind::DefGlobal { global_index, .. } => {
// Phase 4: DCE for Globals
// A global can be removed if it was only used for inlining in this script.
!sub.used_globals.contains(global_index) && self.is_pure(&e)
}
_ => self.is_pure(&e),
};
if !removable {
filtered.push(e);
}
}
new_exprs = filtered;
}
if new_exprs.is_empty() {
return Node { identity: node.identity, kind: BoundKind::Nop, ty: StaticType::Void };
} else if new_exprs.len() == 1 {
@@ -250,13 +282,23 @@ impl Optimizer {
mapping.push(None);
} else {
mapping.push(Some(new_upvalues.len() as u32));
new_upvalues.push(*capture_addr);
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);
}
}
let params = Rc::new(self.visit_node(params.as_ref().clone(), &mut SubstitutionMap::new()));
let body_node = self.visit_node((*body).clone(), &mut next_inner_subs);
// 1. Visit parameters to determine their new slots
let params_node = self.visit_node(params.as_ref().clone(), &mut SubstitutionMap::new());
// 2. IMPORTANT: Pre-register parameter slots in the inner substitution map
// so they are mapped 1:1 and don't get reassigned to higher slots in the body.
self.collect_parameter_slots(&params_node, &mut next_inner_subs);
let body_node = self.visit_node((*body).clone(), &mut next_inner_subs);
let reindexed_body = if new_upvalues.len() != original_upvalues.len() {
sub.reindex_upvalues(body_node, &mapping)
} else {
@@ -264,7 +306,7 @@ impl Optimizer {
};
(BoundKind::Lambda {
params,
params: Rc::new(params_node),
upvalues: new_upvalues,
body: Rc::new(reindexed_body),
positional_count
@@ -281,45 +323,30 @@ impl Optimizer {
} else {
sub.locals.remove(&slot);
}
(BoundKind::DefLocal { name, slot, value, captured_by }, node.ty)
if self.is_pure(&value) {
sub.pure_slots.insert(slot);
} else {
sub.pure_slots.remove(&slot);
}
let new_slot = sub.map_slot(slot);
(BoundKind::DefLocal { name, slot: new_slot, value, captured_by }, node.ty)
},
BoundKind::DefGlobal { name, global_index, value } => {
let value = Box::new(self.visit_node(*value, sub));
// Phase 4: Track script-local global definitions for inlining
if let BoundKind::Constant(val) = &value.kind {
sub.add_global(global_index, val.clone());
}
// Purity Inference: Globals are pure if their value is pure
if self.is_pure(&value) {
if let Some(purity_rc) = &self.global_purity {
purity_rc.borrow_mut().insert(global_index, true);
}
}
(BoundKind::DefGlobal { name, global_index, value }, node.ty)
},
BoundKind::Set { addr, value } => {
let value = Box::new(self.visit_node(*value, sub));
match addr {
Address::Local(slot) => {
if let BoundKind::Constant(val) = &value.kind {
sub.add_local(slot, val.clone());
} else {
sub.locals.remove(&slot);
}
}
Address::Global(idx) => {
if let BoundKind::Constant(val) = &value.kind {
sub.add_global(idx, val.clone());
} else {
sub.globals.remove(&idx);
}
}
_ => {}
}
(BoundKind::Set { addr, value }, node.ty)
},
BoundKind::Tuple { elements } => {
let elements = elements.into_iter().map(|e| self.visit_node(e, sub)).collect();
(BoundKind::Tuple { elements }, node.ty)
@@ -338,11 +365,24 @@ impl Optimizer {
Node { identity: node.identity, kind: new_kind, ty: new_ty }
}
fn collect_parameter_slots(&self, node: &TypedNode, sub: &mut SubstitutionMap) {
match &node.kind {
BoundKind::Parameter { slot, .. } => {
// Identity mapping for parameters to keep them in their reserved frame slots
sub.slot_mapping.insert(*slot, *slot);
sub.next_slot = sub.next_slot.max(*slot + 1);
}
BoundKind::Tuple { elements } => {
for el in elements { self.collect_parameter_slots(el, sub); }
}
_ => {}
}
}
fn is_inlinable_value(&self, val: &Value) -> bool {
match val {
Value::Int(_) | Value::Float(_) | Value::Bool(_) | Value::Text(_) | Value::Keyword(_) | Value::DateTime(_) => true,
Value::Object(obj) => {
// Inline Closures ONLY if they are stateless (cracked)
if let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
closure.upvalues.is_empty()
} else {
@@ -365,7 +405,7 @@ impl Optimizer {
false
}
}
_ => true, // Locals and Upvalues are pure in terms of side-effects
_ => true,
}
}
BoundKind::Tuple { elements } => elements.iter().all(|e| self.is_pure(e)),
@@ -376,9 +416,7 @@ impl Optimizer {
BoundKind::Block { exprs } => exprs.iter().all(|e| self.is_pure(e)),
BoundKind::Expansion { bound_expanded, .. } => self.is_pure(bound_expanded),
BoundKind::DefLocal { value, .. } | BoundKind::Set { addr: Address::Local(_), value } => self.is_pure(value),
BoundKind::Call { callee, args } => {
// A call is pure if the callee is a known pure function and args are pure
let callee_pure = match &callee.kind {
BoundKind::Get { addr: Address::Global(idx), .. } => {
if let Some(purity_rc) = &self.global_purity {
@@ -387,7 +425,7 @@ impl Optimizer {
false
}
}
BoundKind::Constant(Value::Function(_)) => true, // Native fns are checked at registry
BoundKind::Constant(Value::Function(_)) => true,
BoundKind::Constant(Value::Object(obj)) => {
if let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
self.is_pure(&closure.function_node)
@@ -399,38 +437,24 @@ impl Optimizer {
};
callee_pure && self.is_pure(args)
}
_ => false, // DefGlobal and Set(Global) are impure
_ => false,
}
}
fn try_fold_pure(&self, callee: &TypedNode, args: &TypedNode) -> Option<TypedNode> {
// 1. Check if the whole call is pure
// Create a temporary call node to check purity
let temp_call = Node {
identity: callee.identity.clone(),
ty: StaticType::Any,
kind: BoundKind::Call { callee: Box::new(callee.clone()), args: Box::new(args.clone()) }
};
if !self.is_pure(&temp_call) {
return None;
}
// 2. Extract constant arguments
if !self.is_pure(&temp_call) { return None; }
let mut arg_nodes = Vec::new();
self.flatten_typed_tuple(args, &mut arg_nodes);
let mut arg_values = Vec::with_capacity(arg_nodes.len());
for node in arg_nodes {
if let BoundKind::Constant(val) = node.kind {
arg_values.push(val);
} else {
return None; // All arguments must be constants for folding
}
if let BoundKind::Constant(val) = node.kind { arg_values.push(val); }
else { return None; }
}
// 3. Resolve the function to execute
let func_val = match &callee.kind {
BoundKind::Get { addr: Address::Global(idx), .. } => {
self.globals.as_ref()?.borrow().get(*idx as usize)?.clone()
@@ -438,26 +462,10 @@ impl Optimizer {
BoundKind::Constant(val) => val.clone(),
_ => return None,
};
// 4. Execute
let result = match func_val {
Value::Function(f) => f(arg_values),
Value::Object(obj) => {
if let Some(_closure) = obj.as_any().downcast_ref::<Closure>() {
// For closures, we'd need a VM to execute.
// However, if it's pure and we are in the Optimizer,
// we might want to avoid full VM spin-up here if possible,
// or use a shared VM instance.
// For now, we only fold Native Functions.
return None;
} else {
return None;
}
}
_ => return None,
};
// 5. Return new constant node
Some(Node {
identity: callee.identity.clone(),
ty: result.static_type(),
@@ -466,10 +474,6 @@ impl Optimizer {
}
fn try_beta_reduce(&self, params: &TypedNode, args: &TypedNode, body: TypedNode) -> Option<TypedNode> {
if self.contains_def_local(&body) {
return None;
}
let mut sub = SubstitutionMap::new();
let mut arg_vals = Vec::new();
self.flatten_typed_tuple(args, &mut arg_vals);
@@ -477,22 +481,15 @@ impl Optimizer {
let mut slot_index = 0;
self.map_params_to_args(params, &arg_vals, &mut slot_index, &mut sub);
if sub.locals.len() < arg_vals.len() {
return None;
}
if sub.locals.is_empty() && sub.upvalues.is_empty() && !arg_vals.is_empty() {
return None;
}
if sub.locals.len() < arg_vals.len() { return None; }
if sub.locals.is_empty() && sub.upvalues.is_empty() && !arg_vals.is_empty() { return None; }
Some(self.visit_node(body, &mut sub))
}
fn flatten_typed_tuple(&self, node: &TypedNode, into: &mut Vec<TypedNode>) {
if let BoundKind::Tuple { elements } = &node.kind {
for el in elements {
self.flatten_typed_tuple(el, into);
}
for el in elements { self.flatten_typed_tuple(el, into); }
} else if !matches!(node.kind, BoundKind::Nop) {
into.push(node.clone());
}
@@ -505,47 +502,70 @@ impl Optimizer {
&& let BoundKind::Constant(val) = &arg.kind {
sub.add_local(*slot, val.clone());
}
// Parameters in beta-reduction are also identity mapped if not inlined
sub.slot_mapping.insert(*slot, *slot);
sub.next_slot = sub.next_slot.max(*slot + 1);
*offset += 1;
}
BoundKind::Tuple { elements } => {
for el in elements {
self.map_params_to_args(el, args, offset, sub);
}
for el in elements { self.map_params_to_args(el, args, offset, sub); }
}
_ => {}
}
}
fn contains_def_local(&self, node: &TypedNode) -> bool {
fn collect_local_usage(&self, node: &TypedNode, used: &mut HashSet<u32>) {
match &node.kind {
BoundKind::DefLocal { .. } => true,
BoundKind::If { cond, then_br, else_br } => {
self.contains_def_local(cond) || self.contains_def_local(then_br) || else_br.as_ref().is_some_and(|e| self.contains_def_local(e))
BoundKind::Get { addr: Address::Local(slot), .. } | BoundKind::Set { addr: Address::Local(slot), .. } => {
used.insert(*slot);
}
BoundKind::Block { exprs } => exprs.iter().any(|e| self.contains_def_local(e)),
BoundKind::Call { callee, args } => self.contains_def_local(callee) || self.contains_def_local(args),
BoundKind::Lambda { .. } => false,
BoundKind::Tuple { elements } => elements.iter().any(|e| self.contains_def_local(e)),
BoundKind::Record { fields } => fields.iter().any(|(k, v)| self.contains_def_local(k) || self.contains_def_local(v)),
BoundKind::Set { value, .. } => self.contains_def_local(value),
_ => false,
BoundKind::Lambda { upvalues, .. } => {
for addr in upvalues {
if let Address::Local(slot) = addr { used.insert(*slot); }
}
}
BoundKind::Block { exprs } => {
for e in exprs { self.collect_local_usage(e, used); }
}
BoundKind::If { cond, then_br, else_br } => {
self.collect_local_usage(cond, used);
self.collect_local_usage(then_br, used);
if let Some(e) = else_br { self.collect_local_usage(e, used); }
}
BoundKind::Call { callee, args } => {
self.collect_local_usage(callee, used);
self.collect_local_usage(args, used);
}
BoundKind::Tuple { elements } => {
for e in elements { self.collect_local_usage(e, used); }
}
BoundKind::Record { fields } => {
for (k, v) in fields {
self.collect_local_usage(k, used);
self.collect_local_usage(v, used);
}
}
BoundKind::DefLocal { value, .. } | BoundKind::Set { value, .. } => {
self.collect_local_usage(value, used);
}
BoundKind::Expansion { bound_expanded, .. } => {
self.collect_local_usage(bound_expanded, used);
}
_ => {}
}
}
}
/// Helper for transitively inlining values and cleaning up capture lists.
struct SubstitutionMap {
/// Mapping Address::Local(slot) -> Value
locals: HashMap<u32, Value>,
/// Mapping Address::Upvalue(idx) -> Value
upvalues: HashMap<u32, Value>,
/// Mapping Address::Global(idx) -> Value (Script-local substitution)
globals: HashMap<u32, Value>,
/// Tracking which local slots are actually used (not inlined)
slot_mapping: HashMap<u32, u32>,
next_slot: u32,
used_slots: HashSet<u32>,
/// Tracking which global indices are used
pure_slots: HashSet<u32>,
used_globals: HashSet<u32>,
/// Tracking which local slots are captured by lambdas
captured_slots: HashSet<u32>,
}
@@ -555,12 +575,25 @@ impl SubstitutionMap {
locals: HashMap::new(),
upvalues: HashMap::new(),
globals: HashMap::new(),
slot_mapping: HashMap::new(),
next_slot: 0,
used_slots: HashSet::new(),
pure_slots: HashSet::new(),
used_globals: HashSet::new(),
captured_slots: HashSet::new(),
}
}
fn map_slot(&mut self, old_slot: u32) -> u32 {
if let Some(&new_slot) = self.slot_mapping.get(&old_slot) {
return new_slot;
}
let new_slot = self.next_slot;
self.slot_mapping.insert(old_slot, new_slot);
self.next_slot += 1;
new_slot
}
fn add_local(&mut self, slot: u32, val: Value) {
self.locals.insert(slot, val);
}
@@ -573,7 +606,6 @@ impl SubstitutionMap {
self.globals.insert(idx, val);
}
/// Re-maps Get(Upvalue(old_idx)) to Get(Upvalue(new_idx)) based on a mapping table.
fn reindex_upvalues(&self, node: TypedNode, mapping: &[Option<u32>]) -> TypedNode {
let (new_kind, new_ty) = match node.kind {
BoundKind::Get { addr: Address::Upvalue(idx), name } => {
@@ -586,22 +618,18 @@ impl SubstitutionMap {
(BoundKind::Get { addr: Address::Upvalue(idx), name }, node.ty)
}
},
BoundKind::Lambda { params, upvalues, body, positional_count } => {
let mut next_upvalues = Vec::new();
for addr in upvalues {
if let Address::Upvalue(idx) = addr
&& let Some(res) = mapping.get(idx as usize) {
if let Some(new_idx) = res {
next_upvalues.push(Address::Upvalue(*new_idx));
}
if let Some(new_idx) = res { next_upvalues.push(Address::Upvalue(*new_idx)); }
continue;
}
next_upvalues.push(addr);
}
(BoundKind::Lambda { params, upvalues: next_upvalues, body, positional_count }, node.ty)
},
BoundKind::If { cond, then_br, else_br } => {
let cond = Box::new(self.reindex_upvalues(*cond, mapping));
let then_br = Box::new(self.reindex_upvalues(*then_br, mapping));
@@ -639,7 +667,6 @@ impl SubstitutionMap {
},
k => (k, node.ty),
};
Node { identity: node.identity, kind: new_kind, ty: new_ty }
}
}
@@ -662,19 +689,12 @@ mod tests {
#[test]
fn test_opt_folding_complex() {
// Test float folding
let dump_float = get_optimized_dump("(+ 1.5 2.5)", 2);
assert!(dump_float.contains("Constant: 4"));
// Test text folding
let dump_text = get_optimized_dump("(+ \"hello \" \"world\")", 2);
assert!(dump_text.contains("Constant: \"hello world\""));
// Test date folding
let dump_date = get_optimized_dump("(date \"2023-01-01\")", 2);
assert!(dump_date.contains("Constant: #2023-01-01 00:00:00#"));
// Test comparison folding
let dump_cmp = get_optimized_dump("(> 10 5)", 2);
assert!(dump_cmp.contains("Constant: true"));
}
src/ast/environment.rs
+1 -1
View File
@@ -91,7 +91,7 @@ impl Environment {
function_registry: Rc::new(RefCell::new(HashMap::new())),
monomorph_cache: Rc::new(RefCell::new(HashMap::new())),
debug_mode: false,
optimization_level: 0,
optimization_level: 2,
};
env.register_stdlib();
env