diff --git a/src/ast/compiler/optimizer/engine.rs b/src/ast/compiler/optimizer/engine.rs index 6a9661c..16adbfb 100644 --- a/src/ast/compiler/optimizer/engine.rs +++ b/src/ast/compiler/optimizer/engine.rs @@ -1,5 +1,5 @@ use crate::ast::compiler::bound_nodes::{ - Address, AnalyzedNode, BoundKind, GlobalIdx, LocalSlot, UpvalueIdx, + Address, AnalyzedNode, BoundKind, GlobalIdx, UpvalueIdx, }; use crate::ast::nodes::Node; use crate::ast::types::{Purity, Value}; @@ -10,6 +10,7 @@ use std::rc::Rc; use super::substitution_map::{SubstitutionMap, UsageInfo}; use super::folder::Folder; +use super::inliner::Inliner; pub struct Optimizer { pub enabled: bool, @@ -101,12 +102,13 @@ impl Optimizer { path: &mut PathTracker, ) -> AnalyzedNode { let folder = Folder::new(&self.globals); + let inliner = Inliner::new(&self.globals, &self.global_purity); 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) + && inliner.is_inlinable_value(val, addr) { return folder.make_constant_node(val.clone(), &node); } @@ -122,7 +124,7 @@ impl Optimizer { { let globals = globals_rc.borrow(); if let Some(val) = globals.get(idx.0 as usize) - && self.is_inlinable_value(val, addr) + && inliner.is_inlinable_value(val, addr) { return folder.make_constant_node(val.clone(), &node); } @@ -215,13 +217,15 @@ impl Optimizer { && 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, - ); + let mut inner_sub = SubstitutionMap::new(); + let collapsed = if inliner + .prepare_beta_reduction(params, &args, body, &mut inner_sub) + .is_some() + { + Some(self.visit_node((**body).clone(), &mut inner_sub, path)) + } else { + None + }; path.inlining_depth -= 1; path.exit_lambda(&callee.identity); @@ -253,13 +257,14 @@ impl Optimizer { 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, - ); + let collapsed = if inliner + .prepare_beta_reduction(params, &args, body, &mut inner_sub) + .is_some() + { + Some(self.visit_node((**body).clone(), &mut inner_sub, path)) + } else { + None + }; path.inlining_depth -= 1; path.inlining_stack.remove(idx); @@ -292,13 +297,19 @@ impl Optimizer { path, ); - let collapsed = self.try_beta_reduce_with_sub( - &closure.parameter_node, - &args, - inlined_body, - &mut closure_sub, - path, - ); + let collapsed = if inliner + .prepare_beta_reduction( + &closure.parameter_node, + &args, + &inlined_body, + &mut closure_sub, + ) + .is_some() + { + Some(self.visit_node(inlined_body, &mut closure_sub, path)) + } else { + None + }; path.inlining_depth -= 1; path.exit_lambda(&closure.function_node.identity); @@ -468,7 +479,7 @@ impl Optimizer { 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); + inliner.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() { @@ -549,213 +560,4 @@ impl Optimizer { 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.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 try_beta_reduce_with_sub( - &self, - params: &AnalyzedNode, - args: &AnalyzedNode, - body: AnalyzedNode, - sub: &mut SubstitutionMap, - path: &mut PathTracker, - ) -> Option { - 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) { - 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) { - 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); - } - } - _ => {} - } - } } - diff --git a/src/ast/compiler/optimizer/folder.rs b/src/ast/compiler/optimizer/folder.rs index 0744b02..59a7463 100644 --- a/src/ast/compiler/optimizer/folder.rs +++ b/src/ast/compiler/optimizer/folder.rs @@ -4,6 +4,8 @@ use crate::ast::types::{Purity, StaticType, Value}; use std::cell::RefCell; use std::rc::Rc; +use super::utils::flatten_tuple; + pub struct Folder<'a> { pub globals: &'a Option>>>, } @@ -54,7 +56,7 @@ impl<'a> Folder<'a> { } let mut arg_nodes = Vec::new(); - self.flatten_tuple(args.clone(), &mut arg_nodes); + 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 { @@ -77,21 +79,5 @@ impl<'a> Folder<'a> { }; Some(self.make_constant_node(result, callee)) } - - fn flatten_tuple(&self, node: AnalyzedNode, into: &mut Vec) { - 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), - } - } } + diff --git a/src/ast/compiler/optimizer/inliner.rs b/src/ast/compiler/optimizer/inliner.rs new file mode 100644 index 0000000..16280ac --- /dev/null +++ b/src/ast/compiler/optimizer/inliner.rs @@ -0,0 +1,207 @@ +use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, GlobalIdx, LocalSlot}; +use crate::ast::types::{Purity, 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}; +use super::utils::flatten_tuple; + +pub struct Inliner<'a> { + pub globals: &'a Option>>>, + pub global_purity: &'a Option>>>, +} + +impl<'a> Inliner<'a> { + pub fn new( + globals: &'a Option>>>, + global_purity: &'a Option>>>, + ) -> Self { + Self { + globals, + global_purity, + } + } + + pub fn is_inlinable_value(&self, val: &Value, addr: Address) -> bool { + 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.upvalues.is_empty() && !closure.function_node.ty.is_recursive + } else { + false + } + } + _ => false, + }; + + if !type_ok { + return false; + } + + 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; + } + return false; + } + + true + } + + pub fn prepare_beta_reduction( + &self, + params: &AnalyzedNode, + args: &AnalyzedNode, + body: &AnalyzedNode, + sub: &mut SubstitutionMap, + ) -> Option<()> { + let mut body_usage = UsageInfo::default(); + body_usage.collect(body); + + let mut arg_vals = Vec::new(); + 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; + } + + 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(()) + } + + pub 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()); + } + } + + 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 } => { + if let Some(arg) = args.get(*offset) { + let mut sub_args = Vec::new(); + let is_compound = match &arg.kind { + BoundKind::Tuple { .. } | BoundKind::Record { .. } => { + flatten_tuple(arg.clone(), &mut sub_args); + true + } + _ => false, + }; + + if is_compound { + 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; + } + } + + for el in elements { + self.map_params_to_args(el, args, offset, sub, body_usage); + } + } + _ => {} + } + } + + pub fn collect_parameter_slots_set(&self, node: &AnalyzedNode, slots: &mut HashSet) { + 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); + } + } + _ => {} + } + } + + pub 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); + } + } + _ => {} + } + } +} diff --git a/src/ast/compiler/optimizer/mod.rs b/src/ast/compiler/optimizer/mod.rs index 6bf3163..a793532 100644 --- a/src/ast/compiler/optimizer/mod.rs +++ b/src/ast/compiler/optimizer/mod.rs @@ -1,7 +1,10 @@ pub mod engine; pub mod folder; +pub mod inliner; pub mod substitution_map; +pub mod utils; pub use engine::Optimizer; pub use folder::Folder; +pub use inliner::Inliner; pub use substitution_map::{SubstitutionMap, UsageInfo}; diff --git a/src/ast/compiler/optimizer/utils.rs b/src/ast/compiler/optimizer/utils.rs new file mode 100644 index 0000000..d04b5e4 --- /dev/null +++ b/src/ast/compiler/optimizer/utils.rs @@ -0,0 +1,18 @@ +use crate::ast::compiler::bound_nodes::{AnalyzedNode, BoundKind}; + +pub fn flatten_tuple(node: AnalyzedNode, into: &mut Vec) { + match node.kind { + BoundKind::Tuple { elements } => { + for el in elements { + flatten_tuple(el, into); + } + } + BoundKind::Record { fields } => { + for (_, v) in fields { + flatten_tuple(v, into); + } + } + BoundKind::Nop => {} + _ => into.push(node), + } +}