Refactor: Move optimizer related types to dedicated module
This commit reorganizes the optimizer code by creating a new `optimizer` module. The `Optimizer` struct and its related logic remain in `optimizer/optimizer.rs`, while `SubstitutionMap` and `UsageInfo` are moved to `optimizer/substitution_map.rs`. This change improves code organization and makes it easier to manage the optimizer's components.
This commit is contained in:
Michael Schimmel
@@ -0,0 +1,5 @@
|
||||
pub mod optimizer;
|
||||
pub mod substitution_map;
|
||||
|
||||
pub use optimizer::Optimizer;
|
||||
pub use substitution_map::{SubstitutionMap, UsageInfo};
|
||||
@@ -0,0 +1,824 @@
|
||||
use crate::ast::compiler::bound_nodes::{
|
||||
Address, AnalyzedNode, BoundKind, GlobalIdx, LocalSlot, NodeMetrics, UpvalueIdx,
|
||||
};
|
||||
use crate::ast::nodes::Node;
|
||||
use crate::ast::types::{Purity, StaticType, Value};
|
||||
use crate::ast::vm::Closure;
|
||||
use std::cell::RefCell;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::rc::Rc;
|
||||
|
||||
use super::substitution_map::{SubstitutionMap, UsageInfo};
|
||||
|
||||
pub struct Optimizer {
|
||||
pub enabled: bool,
|
||||
max_passes: usize,
|
||||
pub globals: Option<Rc<RefCell<Vec<Value>>>>,
|
||||
pub global_purity: Option<Rc<RefCell<HashMap<GlobalIdx, Purity>>>>,
|
||||
pub lambda_registry: Option<Rc<RefCell<HashMap<GlobalIdx, AnalyzedNode>>>>,
|
||||
}
|
||||
|
||||
struct PathTracker {
|
||||
inlining_depth: usize,
|
||||
inlining_stack: HashSet<GlobalIdx>,
|
||||
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<GlobalIdx, Purity>>>) -> Self {
|
||||
self.global_purity = Some(purity);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_registry(
|
||||
mut self,
|
||||
registry: Rc<RefCell<HashMap<GlobalIdx, AnalyzedNode>>>,
|
||||
) -> Self {
|
||||
self.lambda_registry = Some(registry);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn optimize(&self, node: AnalyzedNode) -> AnalyzedNode {
|
||||
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: AnalyzedNode,
|
||||
sub: &mut SubstitutionMap,
|
||||
path: &mut PathTracker,
|
||||
) -> AnalyzedNode {
|
||||
let (new_kind, metrics) = match node.kind {
|
||||
BoundKind::Get { addr, ref name } => {
|
||||
if !sub.assigned.contains(&addr) {
|
||||
// 1. Try inlining from current value substitution map (locals/globals/upvalues)
|
||||
if let Some(val) = sub.get_value(&addr)
|
||||
&& self.is_inlinable_value(val, addr)
|
||||
{
|
||||
return self.make_constant_node(val.clone(), &node);
|
||||
}
|
||||
|
||||
// 2. Try inlining from AST substitution map (pure expressions)
|
||||
if let Some(inlined_node) = sub.ast_substitutions.get(&addr) {
|
||||
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.0 as usize)
|
||||
&& self.is_inlinable_value(val, addr)
|
||||
{
|
||||
return self.make_constant_node(val.clone(), &node);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
sub.used.insert(addr);
|
||||
(
|
||||
BoundKind::Get {
|
||||
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,
|
||||
kind,
|
||||
ref value,
|
||||
ref captured_by,
|
||||
} => {
|
||||
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: sub.map_address(addr),
|
||||
kind,
|
||||
value: value_opt,
|
||||
captured_by: captured_by.clone(),
|
||||
},
|
||||
node.ty.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
BoundKind::Call { callee, args } => {
|
||||
let callee = self.visit_node(*callee, sub, path);
|
||||
let args = self.visit_node(*args, sub, path);
|
||||
|
||||
if self.enabled {
|
||||
if let BoundKind::Lambda {
|
||||
params,
|
||||
body,
|
||||
upvalues,
|
||||
positional_count,
|
||||
} = &callee.kind
|
||||
&& upvalues.is_empty()
|
||||
&& positional_count.is_some()
|
||||
&& path.inlining_depth < 5
|
||||
&& !callee.ty.is_recursive
|
||||
&& 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;
|
||||
}
|
||||
}
|
||||
|
||||
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()
|
||||
&& !lambda_node.ty.is_recursive
|
||||
{
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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()
|
||||
|| closure.function_node.ty.purity >= Purity::SideEffectFree)
|
||||
&& !closure.function_node.ty.is_recursive
|
||||
&& path.enter_lambda(&closure.function_node.identity)
|
||||
{
|
||||
let mut closure_sub = SubstitutionMap::new();
|
||||
for (i, cell) in closure.upvalues.iter().enumerate() {
|
||||
closure_sub.add_value(
|
||||
Address::Upvalue(UpvalueIdx(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;
|
||||
path.exit_lambda(&closure.function_node.identity);
|
||||
|
||||
if let Some(res) = collapsed {
|
||||
return res;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(folded) = self.try_fold_pure(&callee, &args) {
|
||||
return folded;
|
||||
}
|
||||
}
|
||||
|
||||
(
|
||||
BoundKind::Call {
|
||||
callee: Box::new(callee),
|
||||
args: Box::new(args),
|
||||
},
|
||||
node.ty.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
BoundKind::Again { args } => {
|
||||
let args = self.visit_node(*args, sub, path);
|
||||
(
|
||||
BoundKind::Again {
|
||||
args: Box::new(args),
|
||||
},
|
||||
node.ty.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
BoundKind::If {
|
||||
ref cond,
|
||||
ref then_br,
|
||||
ref else_br,
|
||||
} => {
|
||||
let cond_opt = self.visit_node((**cond).clone(), sub, path);
|
||||
if self.enabled
|
||||
&& let BoundKind::Constant(ref val) = cond_opt.kind
|
||||
{
|
||||
if val.is_truthy() {
|
||||
return self.visit_node((**then_br).clone(), sub, path);
|
||||
} else if let Some(else_node) = else_br {
|
||||
return self.visit_node((**else_node).clone(), sub, path);
|
||||
} else {
|
||||
return self.make_nop_node(&node);
|
||||
}
|
||||
}
|
||||
|
||||
let then_br = Box::new(self.visit_node((**then_br).clone(), sub, path));
|
||||
let else_br = else_br
|
||||
.as_ref()
|
||||
.map(|e| Box::new(self.visit_node((**e).clone(), sub, path)));
|
||||
(
|
||||
BoundKind::If {
|
||||
cond: Box::new(cond_opt),
|
||||
then_br,
|
||||
else_br,
|
||||
},
|
||||
node.ty.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
BoundKind::Block { ref exprs } => {
|
||||
let mut info = UsageInfo::default();
|
||||
if !exprs.is_empty() {
|
||||
for e in exprs {
|
||||
info.collect(e);
|
||||
}
|
||||
}
|
||||
|
||||
sub.assigned.extend(info.assigned.iter().cloned());
|
||||
|
||||
let mut new_exprs = Vec::with_capacity(exprs.len());
|
||||
let last_idx = exprs.len().saturating_sub(1);
|
||||
|
||||
for (i, e) in exprs.iter().enumerate() {
|
||||
let is_last = i == last_idx;
|
||||
if self.enabled && !is_last {
|
||||
let removable = match &e.kind {
|
||||
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,
|
||||
};
|
||||
|
||||
if removable {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
let opt = self.visit_node(e.clone(), 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 self.make_nop_node(&node);
|
||||
} else if new_exprs.len() == 1 {
|
||||
return new_exprs.pop().unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
(BoundKind::Block { exprs: new_exprs }, node.ty.clone())
|
||||
}
|
||||
|
||||
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 info = UsageInfo::default();
|
||||
info.collect(&node);
|
||||
|
||||
let mut new_upvalues = Vec::new();
|
||||
let mut mapping = Vec::new();
|
||||
let mut next_inner_subs = SubstitutionMap::new();
|
||||
next_inner_subs.assigned = info.assigned;
|
||||
|
||||
for (old_idx, capture_addr) in original_upvalues.iter().enumerate() {
|
||||
let mut inlined_val = None;
|
||||
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_value(Address::Upvalue(UpvalueIdx(old_idx as u32)), val);
|
||||
mapping.push(None);
|
||||
} else {
|
||||
mapping.push(Some(new_upvalues.len() as u32));
|
||||
new_upvalues.push(sub.map_address(*capture_addr));
|
||||
}
|
||||
}
|
||||
|
||||
let params_node =
|
||||
self.visit_node(params.as_ref().clone(), &mut next_inner_subs, path);
|
||||
self.collect_parameter_slots(¶ms_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.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
BoundKind::Destructure {
|
||||
ref pattern,
|
||||
ref value,
|
||||
} => {
|
||||
let val_opt = Box::new(self.visit_node((**value).clone(), sub, path));
|
||||
let pat_opt = Box::new(self.visit_node((**pattern).clone(), sub, path));
|
||||
|
||||
let mut info = UsageInfo::default();
|
||||
info.collect_pattern(&pat_opt);
|
||||
for addr in info.assigned {
|
||||
sub.remove_value(&addr);
|
||||
}
|
||||
|
||||
(
|
||||
BoundKind::Destructure {
|
||||
pattern: pat_opt,
|
||||
value: val_opt,
|
||||
},
|
||||
node.ty.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
BoundKind::Tuple { elements } => {
|
||||
let elements = elements
|
||||
.into_iter()
|
||||
.map(|e| self.visit_node(e, sub, path))
|
||||
.collect();
|
||||
(BoundKind::Tuple { elements }, node.ty.clone())
|
||||
}
|
||||
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.clone())
|
||||
}
|
||||
BoundKind::Expansion {
|
||||
ref original_call,
|
||||
ref bound_expanded,
|
||||
} => {
|
||||
path.inlining_depth += 1;
|
||||
let bound_expanded =
|
||||
Box::new(self.visit_node((**bound_expanded).clone(), sub, path));
|
||||
path.inlining_depth -= 1;
|
||||
(
|
||||
BoundKind::Expansion {
|
||||
original_call: original_call.clone(),
|
||||
bound_expanded,
|
||||
},
|
||||
node.ty.clone(),
|
||||
)
|
||||
}
|
||||
k => (k, node.ty.clone()),
|
||||
};
|
||||
|
||||
Node {
|
||||
identity: node.identity,
|
||||
kind: new_kind,
|
||||
ty: metrics,
|
||||
}
|
||||
}
|
||||
|
||||
fn collect_parameter_slots(&self, node: &AnalyzedNode, sub: &mut SubstitutionMap) {
|
||||
match &node.kind {
|
||||
BoundKind::Define {
|
||||
addr: Address::Local(slot),
|
||||
..
|
||||
} => {
|
||||
sub.slot_mapping.insert(*slot, *slot);
|
||||
sub.next_slot = sub.next_slot.max(slot.0 + 1);
|
||||
}
|
||||
BoundKind::Tuple { elements } => {
|
||||
for el in elements {
|
||||
self.collect_parameter_slots(el, sub);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn is_inlinable_value(&self, val: &Value, addr: Address) -> bool {
|
||||
// 1. Basic check: is the type itself inlinable (Constants, pure Closures)?
|
||||
let type_ok = 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() && !closure.function_node.ty.is_recursive
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
_ => false,
|
||||
};
|
||||
|
||||
if !type_ok {
|
||||
return false;
|
||||
}
|
||||
|
||||
// 2. For Globals: check if they are actually pure/stable
|
||||
if let Address::Global(idx) = addr {
|
||||
if let Some(purity_rc) = &self.global_purity {
|
||||
let purity = purity_rc
|
||||
.borrow()
|
||||
.get(&idx)
|
||||
.cloned()
|
||||
.unwrap_or(Purity::Impure);
|
||||
return purity >= Purity::Pure;
|
||||
}
|
||||
// If no purity info is available, we assume the global is unstable (safer)
|
||||
return false;
|
||||
}
|
||||
|
||||
true // Locals/Upvalues that reach here and passed the type check are inlinable
|
||||
}
|
||||
|
||||
fn make_constant_node(&self, val: Value, template: &AnalyzedNode) -> AnalyzedNode {
|
||||
let ty = val.static_type();
|
||||
let typed_original = Rc::new(Node {
|
||||
identity: template.identity.clone(),
|
||||
kind: BoundKind::Constant(val.clone()),
|
||||
ty: ty.clone(),
|
||||
});
|
||||
Node {
|
||||
identity: template.identity.clone(),
|
||||
kind: BoundKind::Constant(val),
|
||||
ty: NodeMetrics {
|
||||
original: typed_original,
|
||||
purity: Purity::Pure,
|
||||
is_recursive: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn make_nop_node(&self, template: &AnalyzedNode) -> AnalyzedNode {
|
||||
let typed_original = Rc::new(Node {
|
||||
identity: template.identity.clone(),
|
||||
kind: BoundKind::Nop,
|
||||
ty: StaticType::Void,
|
||||
});
|
||||
Node {
|
||||
identity: template.identity.clone(),
|
||||
kind: BoundKind::Nop,
|
||||
ty: NodeMetrics {
|
||||
original: typed_original,
|
||||
purity: Purity::Pure,
|
||||
is_recursive: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn try_fold_pure(&self, callee: &AnalyzedNode, args: &AnalyzedNode) -> Option<AnalyzedNode> {
|
||||
if callee.ty.purity < Purity::Pure || args.ty.purity < Purity::Pure {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut arg_nodes = Vec::new();
|
||||
self.flatten_tuple(args.clone(), &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.0 as usize)?.clone(),
|
||||
BoundKind::Constant(val) => val.clone(),
|
||||
_ => return None,
|
||||
};
|
||||
let result = match func_val {
|
||||
Value::Function(f) => (f.func)(arg_values),
|
||||
_ => return None,
|
||||
};
|
||||
Some(self.make_constant_node(result, callee))
|
||||
}
|
||||
|
||||
fn try_beta_reduce_with_sub(
|
||||
&self,
|
||||
params: &AnalyzedNode,
|
||||
args: &AnalyzedNode,
|
||||
body: AnalyzedNode,
|
||||
sub: &mut SubstitutionMap,
|
||||
path: &mut PathTracker,
|
||||
) -> Option<AnalyzedNode> {
|
||||
let mut body_usage = UsageInfo::default();
|
||||
body_usage.collect(&body);
|
||||
|
||||
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, &body_usage);
|
||||
|
||||
if slot_index != arg_vals.len() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Verify that all used/assigned parameters have been substituted.
|
||||
let mut param_slots = HashSet::new();
|
||||
self.collect_parameter_slots_set(params, &mut param_slots);
|
||||
|
||||
for slot in param_slots {
|
||||
let addr = Address::Local(slot);
|
||||
if (body_usage.is_used(&addr) || body_usage.is_assigned(&addr))
|
||||
&& sub.get_value(&addr).is_none()
|
||||
&& !sub.ast_substitutions.contains_key(&addr)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
}
|
||||
|
||||
if sub.values.is_empty() && sub.ast_substitutions.is_empty() && !arg_vals.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(self.visit_node(body, sub, path))
|
||||
}
|
||||
|
||||
fn collect_parameter_slots_set(&self, node: &AnalyzedNode, slots: &mut HashSet<LocalSlot>) {
|
||||
match &node.kind {
|
||||
BoundKind::Define {
|
||||
addr: Address::Local(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 } => {
|
||||
for el in elements {
|
||||
self.flatten_tuple(el, into);
|
||||
}
|
||||
}
|
||||
BoundKind::Record { fields } => {
|
||||
for (_, v) in fields {
|
||||
self.flatten_tuple(v, into);
|
||||
}
|
||||
}
|
||||
BoundKind::Nop => {}
|
||||
_ => into.push(node),
|
||||
}
|
||||
}
|
||||
|
||||
fn map_params_to_args(
|
||||
&self,
|
||||
pattern: &AnalyzedNode,
|
||||
args: &[AnalyzedNode],
|
||||
offset: &mut usize,
|
||||
sub: &mut SubstitutionMap,
|
||||
body_usage: &UsageInfo,
|
||||
) {
|
||||
match &pattern.kind {
|
||||
BoundKind::Define { addr, .. } => {
|
||||
if let Some(arg) = args.get(*offset)
|
||||
&& !body_usage.is_assigned(addr)
|
||||
{
|
||||
if let BoundKind::Constant(val) = &arg.kind {
|
||||
sub.add_value(*addr, val.clone());
|
||||
} else if let BoundKind::Lambda { upvalues, .. } = &arg.kind
|
||||
&& upvalues.is_empty()
|
||||
{
|
||||
sub.add_ast_substitution(*addr, arg.clone());
|
||||
} else if let BoundKind::Get {
|
||||
addr: Address::Global(_),
|
||||
..
|
||||
} = &arg.kind
|
||||
{
|
||||
sub.add_ast_substitution(*addr, arg.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// 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.0 + 1);
|
||||
}
|
||||
*offset += 1;
|
||||
}
|
||||
BoundKind::Tuple { elements } => {
|
||||
// RECURSIVE DESTRUCTURING SUPPORT
|
||||
// Check if the current argument at 'offset' is itself a Tuple or Record literal.
|
||||
if let Some(arg) = args.get(*offset) {
|
||||
let mut sub_args = Vec::new();
|
||||
let is_compound = match &arg.kind {
|
||||
BoundKind::Tuple { .. } | BoundKind::Record { .. } => {
|
||||
self.flatten_tuple(arg.clone(), &mut sub_args);
|
||||
true
|
||||
}
|
||||
_ => false,
|
||||
};
|
||||
|
||||
if is_compound {
|
||||
// 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,
|
||||
body_usage,
|
||||
);
|
||||
}
|
||||
*offset += 1;
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: Continue flat matching (original behavior)
|
||||
for el in elements {
|
||||
self.map_params_to_args(el, args, offset, sub, body_usage);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,321 @@
|
||||
use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, LocalSlot, UpvalueIdx};
|
||||
use crate::ast::nodes::Node;
|
||||
use crate::ast::types::{Identity, Value};
|
||||
use crate::ast::vm::Closure;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct UsageInfo {
|
||||
pub used: HashSet<Address>,
|
||||
pub assigned: HashSet<Address>,
|
||||
pub used_identities: HashSet<Identity>,
|
||||
}
|
||||
|
||||
impl UsageInfo {
|
||||
pub fn is_assigned(&self, addr: &Address) -> bool {
|
||||
self.assigned.contains(addr)
|
||||
}
|
||||
|
||||
pub fn is_used(&self, addr: &Address) -> bool {
|
||||
self.used.contains(addr)
|
||||
}
|
||||
|
||||
pub fn collect(&mut self, node: &AnalyzedNode) {
|
||||
match &node.kind {
|
||||
BoundKind::Destructure { pattern, value } => {
|
||||
self.collect_pattern(pattern);
|
||||
self.collect(value);
|
||||
}
|
||||
BoundKind::Constant(v) => {
|
||||
if let Value::Object(obj) = v
|
||||
&& let Some(closure) = obj.as_any().downcast_ref::<Closure>()
|
||||
{
|
||||
self.used_identities
|
||||
.insert(closure.function_node.identity.clone());
|
||||
self.collect(&closure.function_node);
|
||||
}
|
||||
}
|
||||
BoundKind::Get { addr, .. } => {
|
||||
self.used.insert(*addr);
|
||||
}
|
||||
BoundKind::Set { addr, value } => {
|
||||
self.assigned.insert(*addr);
|
||||
self.collect(value);
|
||||
}
|
||||
BoundKind::Lambda {
|
||||
params,
|
||||
body,
|
||||
upvalues,
|
||||
..
|
||||
} => {
|
||||
self.used_identities.insert(node.identity.clone());
|
||||
|
||||
let mut inner_info = UsageInfo::default();
|
||||
inner_info.collect(params);
|
||||
inner_info.collect(body);
|
||||
|
||||
// Propagate globals and identities
|
||||
for addr in &inner_info.used {
|
||||
if let Address::Global(_) = addr {
|
||||
self.used.insert(*addr);
|
||||
}
|
||||
}
|
||||
for addr in &inner_info.assigned {
|
||||
if let Address::Global(_) = addr {
|
||||
self.assigned.insert(*addr);
|
||||
}
|
||||
}
|
||||
self.used_identities.extend(inner_info.used_identities);
|
||||
|
||||
// Map used upvalues to parent scope
|
||||
for addr in &inner_info.used {
|
||||
if let Address::Upvalue(idx) = addr
|
||||
&& let Some(parent_addr) = upvalues.get(idx.0 as usize)
|
||||
{
|
||||
self.used.insert(*parent_addr);
|
||||
}
|
||||
}
|
||||
// Map assigned upvalues to parent scope
|
||||
for addr in &inner_info.assigned {
|
||||
if let Address::Upvalue(idx) = addr
|
||||
&& let Some(parent_addr) = upvalues.get(idx.0 as usize)
|
||||
{
|
||||
self.assigned.insert(*parent_addr);
|
||||
}
|
||||
}
|
||||
}
|
||||
BoundKind::Block { exprs } => {
|
||||
for e in exprs {
|
||||
self.collect(e);
|
||||
}
|
||||
}
|
||||
BoundKind::If {
|
||||
cond,
|
||||
then_br,
|
||||
else_br,
|
||||
} => {
|
||||
self.collect(cond);
|
||||
self.collect(then_br);
|
||||
if let Some(e) = else_br {
|
||||
self.collect(e);
|
||||
}
|
||||
}
|
||||
BoundKind::Call { callee, args } => {
|
||||
self.collect(callee);
|
||||
self.collect(args);
|
||||
}
|
||||
BoundKind::Tuple { elements } => {
|
||||
for e in elements {
|
||||
self.collect(e);
|
||||
}
|
||||
}
|
||||
BoundKind::Record { fields } => {
|
||||
for (k, v) in fields {
|
||||
self.collect(k);
|
||||
self.collect(v);
|
||||
}
|
||||
}
|
||||
BoundKind::Define { value, .. } => {
|
||||
self.collect(value);
|
||||
}
|
||||
BoundKind::Expansion { bound_expanded, .. } => {
|
||||
self.collect(bound_expanded);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn collect_pattern(&mut self, node: &AnalyzedNode) {
|
||||
match &node.kind {
|
||||
BoundKind::Define { addr, .. } => {
|
||||
self.assigned.insert(*addr);
|
||||
}
|
||||
BoundKind::Set { addr, .. } => {
|
||||
self.assigned.insert(*addr);
|
||||
}
|
||||
BoundKind::Tuple { elements } => {
|
||||
for el in elements {
|
||||
self.collect_pattern(el);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SubstitutionMap {
|
||||
pub values: HashMap<Address, Value>,
|
||||
pub ast_substitutions: HashMap<Address, AnalyzedNode>,
|
||||
pub slot_mapping: HashMap<LocalSlot, LocalSlot>,
|
||||
pub assigned: HashSet<Address>,
|
||||
pub next_slot: u32,
|
||||
pub used: HashSet<Address>,
|
||||
pub captured_slots: HashSet<LocalSlot>,
|
||||
}
|
||||
|
||||
impl SubstitutionMap {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
values: HashMap::new(),
|
||||
ast_substitutions: HashMap::new(),
|
||||
slot_mapping: HashMap::new(),
|
||||
assigned: HashSet::new(),
|
||||
next_slot: 0,
|
||||
used: HashSet::new(),
|
||||
captured_slots: HashSet::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_ast_substitution(&mut self, addr: Address, node: AnalyzedNode) {
|
||||
self.ast_substitutions.insert(addr, node);
|
||||
}
|
||||
|
||||
pub fn map_slot(&mut self, old_slot: LocalSlot) -> LocalSlot {
|
||||
if let Some(&new_slot) = self.slot_mapping.get(&old_slot) {
|
||||
return new_slot;
|
||||
}
|
||||
let new_slot = LocalSlot(self.next_slot);
|
||||
self.slot_mapping.insert(old_slot, new_slot);
|
||||
self.next_slot += 1;
|
||||
new_slot
|
||||
}
|
||||
|
||||
pub fn map_address(&mut self, addr: Address) -> Address {
|
||||
match addr {
|
||||
Address::Local(slot) => Address::Local(self.map_slot(slot)),
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_value(&mut self, addr: Address, val: Value) {
|
||||
self.values.insert(addr, val);
|
||||
}
|
||||
|
||||
pub fn get_value(&self, addr: &Address) -> Option<&Value> {
|
||||
self.values.get(addr)
|
||||
}
|
||||
|
||||
pub fn remove_value(&mut self, addr: &Address) {
|
||||
self.values.remove(addr);
|
||||
}
|
||||
|
||||
fn reindex_addr(&self, addr: Address, mapping: &[Option<u32>]) -> Address {
|
||||
if let Address::Upvalue(idx) = addr
|
||||
&& let Some(res) = mapping.get(idx.0 as usize)
|
||||
&& let Some(new_idx) = res
|
||||
{
|
||||
Address::Upvalue(UpvalueIdx(*new_idx))
|
||||
} else {
|
||||
addr
|
||||
}
|
||||
}
|
||||
|
||||
pub fn reindex_upvalues(&self, node: AnalyzedNode, mapping: &[Option<u32>]) -> AnalyzedNode {
|
||||
let (new_kind, metrics) = match node.kind {
|
||||
BoundKind::Get { addr, name } => (
|
||||
BoundKind::Get {
|
||||
addr: self.reindex_addr(addr, mapping),
|
||||
name,
|
||||
},
|
||||
node.ty.clone(),
|
||||
),
|
||||
BoundKind::Lambda {
|
||||
params,
|
||||
upvalues,
|
||||
body,
|
||||
positional_count,
|
||||
} => {
|
||||
let mut next_upvalues = Vec::new();
|
||||
for addr in upvalues {
|
||||
next_upvalues.push(self.reindex_addr(addr, mapping));
|
||||
}
|
||||
(
|
||||
BoundKind::Lambda {
|
||||
params,
|
||||
upvalues: next_upvalues,
|
||||
body,
|
||||
positional_count,
|
||||
},
|
||||
node.ty.clone(),
|
||||
)
|
||||
}
|
||||
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.clone(),
|
||||
)
|
||||
}
|
||||
BoundKind::Block { exprs } => {
|
||||
let exprs = exprs
|
||||
.into_iter()
|
||||
.map(|e| self.reindex_upvalues(e, mapping))
|
||||
.collect();
|
||||
(BoundKind::Block { exprs }, node.ty.clone())
|
||||
}
|
||||
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.clone())
|
||||
}
|
||||
BoundKind::Define {
|
||||
name,
|
||||
addr,
|
||||
kind,
|
||||
value,
|
||||
captured_by,
|
||||
} => {
|
||||
let value = Box::new(self.reindex_upvalues(*value, mapping));
|
||||
(
|
||||
BoundKind::Define {
|
||||
name,
|
||||
addr,
|
||||
kind,
|
||||
value,
|
||||
captured_by,
|
||||
},
|
||||
node.ty.clone(),
|
||||
)
|
||||
}
|
||||
BoundKind::Set { addr, value } => {
|
||||
let value = Box::new(self.reindex_upvalues(*value, mapping));
|
||||
(BoundKind::Set { addr, value }, node.ty.clone())
|
||||
}
|
||||
BoundKind::Tuple { elements } => {
|
||||
let elements = elements
|
||||
.into_iter()
|
||||
.map(|e| self.reindex_upvalues(e, mapping))
|
||||
.collect();
|
||||
(BoundKind::Tuple { elements }, node.ty.clone())
|
||||
}
|
||||
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.clone())
|
||||
}
|
||||
k => (k, node.ty.clone()),
|
||||
};
|
||||
Node {
|
||||
identity: node.identity.clone(),
|
||||
kind: new_kind,
|
||||
ty: metrics,
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user