Files
RustAst/src/ast/compiler/optimizer.rs
T
Michael Schimmel bd506b1b17 feat: Track optimization path to prevent infinite recursion
Introduce a `PathTracker` to manage recursion depth and detect cycles
during optimization. This prevents infinite inlining of recursive
functions.
2026-02-22 10:05:56 +01:00

1440 lines
52 KiB
Rust

use crate::ast::compiler::bound_nodes::{Address, BoundKind, TypedNode};
use crate::ast::nodes::Node;
use crate::ast::types::{StaticType, Value};
use crate::ast::vm::Closure;
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::rc::Rc;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Purity {
Impure, // Has side effects (e.g. Set, Print)
SideEffectFree, // No side effects, but not deterministic (e.g. now(), random())
Pure, // No side effects and deterministic (e.g. sin(x), 1 + 2)
}
/// The Optimizer performs Phase 2 (Cracking), Phase 2.5/2.6 (Aggressive Collapsing & DCE)
/// and Phase 4 (Global Inlining).
pub struct Optimizer {
pub enabled: bool,
max_passes: usize,
pub globals: Option<Rc<RefCell<Vec<Value>>>>,
pub global_purity: Option<Rc<RefCell<HashMap<u32, Purity>>>>,
pub lambda_registry: Option<Rc<RefCell<HashMap<u32, TypedNode>>>>,
}
/// Global state for tracking the optimization path (recursion, depth).
struct PathTracker {
inlining_depth: usize,
/// Stack of global indices to detect recursion in named functions.
inlining_stack: HashSet<u32>,
/// Stack of node identities to detect recursion in anonymous lambdas.
identity_stack: HashSet<crate::ast::types::Identity>,
}
impl PathTracker {
fn new() -> Self {
Self {
inlining_depth: 0,
inlining_stack: HashSet::new(),
identity_stack: HashSet::new(),
}
}
fn enter_lambda(&mut self, identity: &crate::ast::types::Identity) -> bool {
if self.identity_stack.contains(identity) {
return false;
}
self.identity_stack.insert(identity.clone());
true
}
fn exit_lambda(&mut self, identity: &crate::ast::types::Identity) {
self.identity_stack.remove(identity);
}
}
impl Optimizer {
pub fn new(enabled: bool) -> Self {
Self {
enabled,
max_passes: 5,
globals: None,
global_purity: None,
lambda_registry: None,
}
}
pub fn with_globals(mut self, globals: Rc<RefCell<Vec<Value>>>) -> Self {
self.globals = Some(globals);
self
}
pub fn with_purity(mut self, purity: Rc<RefCell<HashMap<u32, Purity>>>) -> Self {
self.global_purity = Some(purity);
self
}
pub fn with_registry(mut self, registry: Rc<RefCell<HashMap<u32, TypedNode>>>) -> Self {
self.lambda_registry = Some(registry);
self
}
pub fn optimize(&self, node: TypedNode) -> TypedNode {
if !self.enabled {
return node;
}
let mut current = node;
for _ in 0..self.max_passes {
let mut sub = SubstitutionMap::new();
let mut path = PathTracker::new();
let next = self.visit_node(current.clone(), &mut sub, &mut path);
if next == current {
break;
}
current = next;
}
current
}
fn visit_node(
&self,
node: TypedNode,
sub: &mut SubstitutionMap,
path: &mut PathTracker,
) -> TypedNode {
let (new_kind, new_ty) = match node.kind {
BoundKind::Get { addr, 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 Node {
identity: node.identity,
ty: val.static_type(),
kind: BoundKind::Constant(val.clone()),
};
}
}
sub.used_slots.insert(slot);
Address::Local(sub.map_slot(slot))
}
Address::Upvalue(idx) => {
if let Some(val) = sub.upvalues.get(&idx) {
return Node {
identity: node.identity,
ty: val.static_type(),
kind: BoundKind::Constant(val.clone()),
};
}
Address::Upvalue(idx)
}
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 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)
{
return Node {
identity: node.identity,
ty: val.static_type(),
kind: BoundKind::Constant(val.clone()),
};
}
}
}
sub.used_globals.insert(idx);
Address::Global(idx)
}
};
(
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, 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,
)
}
BoundKind::Call { callee, args } => {
let callee = self.visit_node(*callee, sub, path);
let args = self.visit_node(*args, sub, path);
if self.enabled {
// Case 1: Beta-Reduction for Lambda Literals
if let BoundKind::Lambda {
params,
body,
upvalues,
positional_count,
} = &callee.kind
&& upvalues.is_empty()
&& positional_count.is_some()
&& path.inlining_depth < 5
{
if path.enter_lambda(&callee.identity) {
path.inlining_depth += 1;
let collapsed = self.try_beta_reduce_with_sub(
params,
&args,
(**body).clone(),
&mut SubstitutionMap::new(),
path,
);
path.inlining_depth -= 1;
path.exit_lambda(&callee.identity);
if let Some(res) = collapsed {
return res;
}
}
}
// Case 1.5: Beta-Reduction for Global Registry functions (from Registry)
if let BoundKind::Get {
addr: Address::Global(idx),
..
} = &callee.kind
&& let Some(registry_rc) = &self.lambda_registry
&& path.inlining_depth < 5
&& !path.inlining_stack.contains(idx)
{
let registry = registry_rc.borrow();
if let Some(lambda_node) = registry.get(idx)
&& let BoundKind::Lambda {
params,
body,
upvalues,
positional_count,
} = &lambda_node.kind
&& 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) {
let mut inner_sub = SubstitutionMap::new();
path.inlining_stack.insert(*idx);
path.inlining_depth += 1;
let collapsed = self.try_beta_reduce_with_sub(
params.as_ref(),
&args,
body.as_ref().clone(),
&mut inner_sub,
path,
);
path.inlining_depth -= 1;
path.inlining_stack.remove(idx);
if let Some(res) = collapsed {
return res;
}
}
}
}
// Case 2: Cracking and Inlining for Constant Closures
if let BoundKind::Constant(Value::Object(ref obj)) = callee.kind
&& path.inlining_depth < 5
&& let Some(closure) = obj.as_any().downcast_ref::<Closure>()
&& (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());
}
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;
if let Some(res) = collapsed {
return res;
}
}
if let Some(folded) = self.try_fold_pure(&callee, &args) {
return folded;
}
}
if self.enabled
&& let BoundKind::Constant(Value::Object(ref obj)) = callee.kind
&& path.inlining_depth < 5
&& let Some(closure) = obj.as_any().downcast_ref::<Closure>()
&& (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());
}
path.inlining_depth += 1;
let inlined_body =
self.visit_node((*closure.function_node).clone(), &mut closure_sub, path);
path.inlining_depth -= 1;
let cracked_lambda = Node {
identity: callee.identity.clone(),
ty: callee.ty.clone(),
kind: BoundKind::Lambda {
params: closure.parameter_node.clone(),
upvalues: vec![],
body: Rc::new(inlined_body),
positional_count: closure.positional_count,
},
};
return Node {
identity: node.identity,
kind: BoundKind::Call {
callee: Box::new(cracked_lambda),
args: Box::new(args),
},
ty: node.ty,
};
}
(
BoundKind::Call {
callee: Box::new(callee),
args: Box::new(args),
},
node.ty,
)
}
BoundKind::If {
cond,
then_br,
else_br,
} => {
let cond = self.visit_node(*cond, sub, path);
if self.enabled
&& let BoundKind::Constant(ref val) = cond.kind
{
if val.is_truthy() {
return self.visit_node(*then_br, sub, path);
} else if let Some(else_node) = else_br {
return self.visit_node(*else_node, sub, path);
} else {
return Node {
identity: node.identity,
kind: BoundKind::Nop,
ty: StaticType::Void,
};
}
}
let then_br = Box::new(self.visit_node(*then_br, sub, path));
let else_br = else_br.map(|e| Box::new(self.visit_node(*e, sub, path)));
(
BoundKind::If {
cond: Box::new(cond),
then_br,
else_br,
},
node.ty,
)
}
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();
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,
);
}
if let Some(last) = exprs.last() {
match &last.kind {
BoundKind::Get {
addr: Address::Local(slot),
..
} => {
used_in_block.insert(*slot);
}
BoundKind::Get {
addr: Address::Global(idx),
..
} => {
used_globals_in_block.insert(*idx);
}
_ => {}
}
}
}
// 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 {
sub.assigned_locals.insert(*slot);
}
for idx in &assigned_globals_in_block {
sub.assigned_globals.insert(*idx);
}
let mut new_exprs = Vec::with_capacity(exprs.len());
let last_idx = exprs.len().saturating_sub(1);
for (i, e) in exprs.into_iter().enumerate() {
let is_last = i == last_idx;
if self.enabled && !is_last {
let removable = match &e.kind {
BoundKind::DefLocal { slot, value, .. } => {
!used_in_block.contains(slot)
&& !sub.captured_slots.contains(slot)
&& self.purity_of(value) >= Purity::SideEffectFree
}
BoundKind::DefGlobal {
global_index,
value,
..
} => {
!used_globals_in_block.contains(global_index)
&& (self.purity_of(value) >= Purity::SideEffectFree
|| matches!(value.kind, BoundKind::Lambda { .. }))
}
BoundKind::Set {
addr: Address::Local(slot),
value,
..
} => {
!used_in_block.contains(slot)
&& !sub.captured_slots.contains(slot)
&& self.purity_of(value) >= Purity::SideEffectFree
}
_ => self.purity_of(&e) >= Purity::SideEffectFree,
};
if removable {
continue;
}
}
let opt = self.visit_node(e, sub, path);
if self.enabled && matches!(opt.kind, BoundKind::Nop) && !is_last {
continue;
}
new_exprs.push(opt);
}
if self.enabled {
if new_exprs.is_empty() {
return Node {
identity: node.identity,
kind: BoundKind::Nop,
ty: StaticType::Void,
};
} else if new_exprs.len() == 1 {
return new_exprs.pop().unwrap();
}
}
let ty = if new_exprs.is_empty() {
StaticType::Void
} else {
new_exprs.last().unwrap().ty.clone()
};
(BoundKind::Block { exprs: new_exprs }, ty)
}
BoundKind::Lambda { .. } => {
let (params, original_upvalues, body, positional_count) =
if let BoundKind::Lambda {
params,
upvalues,
body,
positional_count,
} = &node.kind
{
(params, upvalues, body, positional_count)
} else {
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,
);
// Track captures from outer scope that are assigned in this lambda
for idx in &assigned_globals_in_lambda {
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;
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) = inlined_val {
next_inner_subs.add_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);
}
}
// 1. Visit parameters to determine their new slots
let params_node =
self.visit_node(params.as_ref().clone(), &mut next_inner_subs, path);
// 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.as_ref().clone(), &mut next_inner_subs, path);
let reindexed_body = if new_upvalues.len() != original_upvalues.len() {
sub.reindex_upvalues(body_node, &mapping)
} else {
body_node
};
(
BoundKind::Lambda {
params: Rc::new(params_node),
upvalues: new_upvalues,
body: Rc::new(reindexed_body),
positional_count: *positional_count,
},
node.ty,
)
}
BoundKind::DefLocal {
name,
slot,
value,
captured_by,
} => {
if !captured_by.is_empty() {
sub.captured_slots.insert(slot);
}
let value = Box::new(self.visit_node(*value, sub, path));
if let BoundKind::Constant(val) = &value.kind {
sub.add_local(slot, val.clone());
} else {
sub.locals.remove(&slot);
}
if self.purity_of(&value) == Purity::Pure {
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, path));
if let BoundKind::Constant(val) = &value.kind {
sub.add_global(global_index, val.clone());
}
let p = self.purity_of(&value);
if p > Purity::Impure
&& let Some(purity_rc) = &self.global_purity
{
purity_rc.borrow_mut().insert(global_index, p);
}
(
BoundKind::DefGlobal {
name,
global_index,
value,
},
node.ty,
)
}
BoundKind::Tuple { elements } => {
let elements = elements
.into_iter()
.map(|e| self.visit_node(e, sub, path))
.collect();
(BoundKind::Tuple { elements }, node.ty)
}
BoundKind::Record { fields } => {
let fields = fields
.into_iter()
.map(|(k, v)| {
(
self.visit_node(k, sub, path),
self.visit_node(v, sub, path),
)
})
.collect();
(BoundKind::Record { fields }, node.ty)
}
BoundKind::Expansion {
original_call,
bound_expanded,
} => {
path.inlining_depth += 1;
let bound_expanded = Box::new(self.visit_node(*bound_expanded, sub, path));
path.inlining_depth -= 1;
(
BoundKind::Expansion {
original_call,
bound_expanded,
},
node.ty,
)
}
k => (k, node.ty),
};
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) => {
if let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
closure.upvalues.is_empty()
} else {
false
}
}
_ => false,
}
}
fn purity_of(&self, node: &TypedNode) -> Purity {
match &node.kind {
BoundKind::Constant(_) | BoundKind::Parameter { .. } | BoundKind::Nop => Purity::Pure,
// Defining a lambda is pure; only calling it may have side effects.
BoundKind::Lambda { .. } => Purity::Pure,
BoundKind::Get { addr, .. } => match addr {
Address::Global(idx) => {
if let Some(purity_rc) = &self.global_purity {
*purity_rc.borrow().get(idx).unwrap_or(&Purity::Impure)
} else {
Purity::Impure
}
}
_ => Purity::Pure,
},
BoundKind::Tuple { elements } => elements
.iter()
.map(|e| self.purity_of(e))
.min()
.unwrap_or(Purity::Pure),
BoundKind::Record { fields } => fields
.iter()
.map(|(k, v)| self.purity_of(k).min(self.purity_of(v)))
.min()
.unwrap_or(Purity::Pure),
BoundKind::If {
cond,
then_br,
else_br,
} => {
let p = self.purity_of(cond).min(self.purity_of(then_br));
if let Some(e) = else_br {
p.min(self.purity_of(e))
} else {
p
}
}
BoundKind::Block { exprs } => exprs
.iter()
.map(|e| self.purity_of(e))
.min()
.unwrap_or(Purity::Pure),
BoundKind::Expansion { bound_expanded, .. } => self.purity_of(bound_expanded),
BoundKind::DefLocal { value, .. } => self.purity_of(value),
BoundKind::Set { .. } => Purity::Impure,
BoundKind::Call { callee, args } => {
let callee_purity = match &callee.kind {
BoundKind::Lambda { body, .. } => self.purity_of(body),
BoundKind::Get {
addr: Address::Global(idx),
..
} => {
if let Some(purity_rc) = &self.global_purity {
*purity_rc.borrow().get(idx).unwrap_or(&Purity::Impure)
} else {
Purity::Impure
}
}
BoundKind::Constant(Value::Function(_)) => Purity::Pure,
BoundKind::Constant(Value::Object(obj)) => {
if let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
self.purity_of(&closure.function_node)
} else {
Purity::Impure
}
}
_ => Purity::Impure,
};
callee_purity.min(self.purity_of(args))
}
_ => Purity::Impure,
}
}
fn try_fold_pure(&self, callee: &TypedNode, args: &TypedNode) -> Option<TypedNode> {
let temp_call = Node {
identity: callee.identity.clone(),
ty: StaticType::Any,
kind: BoundKind::Call {
callee: Box::new(callee.clone()),
args: Box::new(args.clone()),
},
};
// Folding requires the function AND arguments to be truly Pure (deterministic).
if self.purity_of(&temp_call) < Purity::Pure {
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;
}
}
let func_val = match &callee.kind {
BoundKind::Get {
addr: Address::Global(idx),
..
} => self.globals.as_ref()?.borrow().get(*idx as usize)?.clone(),
BoundKind::Constant(val) => val.clone(),
_ => return None,
};
let result = match func_val {
Value::Function(f) => f(arg_values),
_ => return None,
};
Some(Node {
identity: callee.identity.clone(),
ty: result.static_type(),
kind: BoundKind::Constant(result),
})
}
fn try_beta_reduce_with_sub(
&self,
params: &TypedNode,
args: &TypedNode,
body: TypedNode,
sub: &mut SubstitutionMap,
path: &mut PathTracker,
) -> Option<TypedNode> {
let mut arg_vals = Vec::new();
self.flatten_typed_tuple(args, &mut arg_vals);
let mut slot_index = 0;
self.map_params_to_args(params, &arg_vals, &mut slot_index, sub);
// Safety: Only proceed if we mapped EXACTLY the number of arguments provided.
// And if we actually have parameters to bind (or if it's a 0-arg call).
if slot_index != arg_vals.len() {
return None;
}
if sub.locals.is_empty()
&& sub.ast_locals.is_empty()
&& sub.upvalues.is_empty()
&& !arg_vals.is_empty()
{
return None;
}
Some(self.visit_node(body, sub, path))
}
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);
}
} else if !matches!(node.kind, BoundKind::Nop) {
into.push(node.clone());
}
}
fn map_params_to_args(
&self,
pattern: &TypedNode,
args: &[TypedNode],
offset: &mut usize,
sub: &mut SubstitutionMap,
) {
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() {
sub.add_ast_local(*slot, arg.clone());
}
} else if let BoundKind::Get {
addr: Address::Global(_),
..
} = &arg.kind
{
sub.add_ast_local(*slot, arg.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);
}
}
_ => {}
}
}
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>,
) {
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,
);
}
}
BoundKind::Get { addr, .. } => match addr {
Address::Local(slot) => {
used_locals.insert(*slot);
}
Address::Global(idx) => {
used_globals.insert(*idx);
}
_ => {}
},
BoundKind::Set { addr, value } => {
match addr {
Address::Local(slot) => {
assigned_locals.insert(*slot);
}
Address::Global(idx) => {
assigned_globals.insert(*idx);
}
_ => {}
}
self.collect_usage(
value,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
}
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::Block { exprs } => {
for e in exprs {
self.collect_usage(
e,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
}
}
BoundKind::If {
cond,
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,
);
if let Some(e) = else_br {
self.collect_usage(
e,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
}
}
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,
);
}
BoundKind::Tuple { elements } => {
for e in elements {
self.collect_usage(
e,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
}
}
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,
);
}
}
BoundKind::DefLocal { value, .. } | BoundKind::DefGlobal { value, .. } => {
self.collect_usage(
value,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
}
BoundKind::Expansion { bound_expanded, .. } => {
self.collect_usage(
bound_expanded,
used_locals,
used_globals,
assigned_locals,
assigned_globals,
);
}
_ => {}
}
}
}
/// Helper for transitively inlining values and cleaning up capture lists.
struct SubstitutionMap {
locals: HashMap<u32, Value>,
upvalues: HashMap<u32, Value>,
globals: HashMap<u32, Value>,
ast_locals: HashMap<u32, TypedNode>,
slot_mapping: HashMap<u32, u32>,
assigned_locals: HashSet<u32>,
assigned_globals: HashSet<u32>,
next_slot: u32,
used_slots: HashSet<u32>,
pure_slots: HashSet<u32>,
used_globals: HashSet<u32>,
captured_slots: HashSet<u32>,
}
impl SubstitutionMap {
fn new() -> Self {
Self {
locals: HashMap::new(),
upvalues: HashMap::new(),
globals: HashMap::new(),
ast_locals: HashMap::new(),
slot_mapping: HashMap::new(),
assigned_locals: HashSet::new(),
assigned_globals: HashSet::new(),
next_slot: 0,
used_slots: HashSet::new(),
pure_slots: HashSet::new(),
used_globals: HashSet::new(),
captured_slots: HashSet::new(),
}
}
fn add_ast_local(&mut self, slot: u32, node: TypedNode) {
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;
}
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);
}
fn add_upvalue(&mut self, idx: u32, val: Value) {
self.upvalues.insert(idx, val);
}
fn add_global(&mut self, idx: u32, val: Value) {
self.globals.insert(idx, val);
}
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,
} => {
if let Some(res) = mapping.get(idx as usize) {
match res {
Some(new_idx) => (
BoundKind::Get {
addr: Address::Upvalue(*new_idx),
name,
},
node.ty,
),
None => (
BoundKind::Get {
addr: Address::Upvalue(idx),
name,
},
node.ty,
),
}
} else {
(
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));
}
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));
let else_br = else_br.map(|e| Box::new(self.reindex_upvalues(*e, mapping)));
(
BoundKind::If {
cond,
then_br,
else_br,
},
node.ty,
)
}
BoundKind::Block { exprs } => {
let exprs = exprs
.into_iter()
.map(|e| self.reindex_upvalues(e, mapping))
.collect();
(BoundKind::Block { exprs }, node.ty)
}
BoundKind::Call { callee, args } => {
let callee = Box::new(self.reindex_upvalues(*callee, mapping));
let args = Box::new(self.reindex_upvalues(*args, mapping));
(BoundKind::Call { callee, args }, node.ty)
}
BoundKind::DefLocal {
name,
slot,
value,
captured_by,
} => {
let value = Box::new(self.reindex_upvalues(*value, mapping));
(
BoundKind::DefLocal {
name,
slot,
value,
captured_by,
},
node.ty,
)
}
BoundKind::DefGlobal {
name,
global_index,
value,
} => {
let value = Box::new(self.reindex_upvalues(*value, mapping));
(
BoundKind::DefGlobal {
name,
global_index,
value,
},
node.ty,
)
}
BoundKind::Set { addr, value } => {
let value = Box::new(self.reindex_upvalues(*value, mapping));
(BoundKind::Set { addr, value }, node.ty)
}
BoundKind::Tuple { elements } => {
let elements = elements
.into_iter()
.map(|e| self.reindex_upvalues(e, mapping))
.collect();
(BoundKind::Tuple { elements }, node.ty)
}
BoundKind::Record { fields } => {
let fields = fields
.into_iter()
.map(|(k, v)| {
(
self.reindex_upvalues(k, mapping),
self.reindex_upvalues(v, mapping),
)
})
.collect();
(BoundKind::Record { fields }, node.ty)
}
k => (k, node.ty),
};
Node {
identity: node.identity,
kind: new_kind,
ty: new_ty,
}
}
}
#[cfg(test)]
mod tests {
use crate::ast::environment::Environment;
fn get_optimized_dump(source: &str, enabled: bool) -> String {
let mut env = Environment::new();
env.optimization = enabled;
env.dump_ast(source)
.expect("Compilation failed during test")
}
#[test]
fn test_opt_folding() {
let dump = get_optimized_dump("(+ 10 20)", true);
insta::assert_snapshot!(dump);
}
#[test]
fn test_opt_folding_complex() {
let dump_float = get_optimized_dump("(+ 1.5 2.5)", true);
assert!(dump_float.contains("Constant: 4"));
let dump_text = get_optimized_dump("(+ \"hello \" \"world\")", true);
assert!(dump_text.contains("Constant: \"hello world\""));
let dump_date = get_optimized_dump("(date \"2023-01-01\")", true);
assert!(dump_date.contains("Constant: #2023-01-01 00:00:00#"));
let dump_cmp = get_optimized_dump("(> 10 5)", true);
assert!(dump_cmp.contains("Constant: true"));
}
#[test]
fn test_opt_local_inlining() {
let source = "(fn [] (do (def x 10) (+ x 5)))";
let dump = get_optimized_dump(source, true);
insta::assert_snapshot!(dump);
}
#[test]
fn test_opt_dce_unused() {
let source = "(fn [] (do (def x 10) 42))";
let dump = get_optimized_dump(source, true);
insta::assert_snapshot!(dump);
}
#[test]
fn test_opt_assignment_safety() {
let source = "(fn [] (do (def x 10) (assign x 20) x))";
let dump = get_optimized_dump(source, true);
insta::assert_snapshot!(dump);
}
#[test]
fn test_opt_lambda_cracking() {
let source = "(fn [] (do (def x 10) (fn [] x)))";
let dump = get_optimized_dump(source, true);
insta::assert_snapshot!(dump);
}
#[test]
fn test_opt_global_unused_dce() {
let source = "(do (def f (fn [x] x)) 42)";
let dump = get_optimized_dump(source, true);
assert!(
!dump.contains("DefGlobal"),
"Unused global helper 'f' should be optimized away. Dump: \n{}",
dump
);
assert!(dump.contains("Constant: 42"));
}
#[test]
fn test_opt_local_recursion_safety() {
// (fn [] (do (def f (fn [x] (f x))) (f 1)))
// Hier sollte der Optimizer 'f' nicht unendlich oft in sich selbst inlinen.
let source = "(fn [] (do (def f (fn [x] (f x))) (f 1)))";
let dump = get_optimized_dump(source, true);
// Der Dump sollte 'Call' auf 'f' (oder Slot-Get) noch enthalten,
// statt 5 Ebenen tief entfaltet zu sein.
assert!(dump.contains("Call"), "Recursive call should remain as a call. Dump: \n{}", dump);
assert!(!dump.contains("- Capturer:"), "Should not be over-optimized. Dump: \n{}", dump);
}
}