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

Refactor: Move inlining logic to Inliner struct

Browse Source 
Introduces the `Inliner` struct to encapsulate logic related to value
and function inlining. This improves code organization and readability
within the optimizer engine.

Key changes include:
- Moving `is_inlinable_value` to the `Inliner` struct.
- Extracting beta-reduction preparation logic into
  `Inliner::prepare_beta_reduction`.
- Moving parameter slot collection to
  `Inliner::collect_parameter_slots`.
- Introducing `flatten_tuple` to a new `utils` module, used by both
  `Folder` and `Inliner`.
This commit is contained in:
Michael Schimmel
2026-02-25 20:07:37 +01:00
parent 7f64e7e6ea
commit cad0c03973
5 changed files with 268 additions and 252 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ast/compiler/optimizer/engine.rs
+36 -234
View File
@@ -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(&params_node, &mut next_inner_subs);
inliner.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() {
@@ -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>() {
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<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);
}
}
_ => {}
}
}
}
src/ast/compiler/optimizer/folder.rs
+4 -18
View File
@@ -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<Rc<RefCell<Vec<Value>>>>,
}
@@ -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<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),
}
}
}
src/ast/compiler/optimizer/inliner.rs
+207
View File
@@ -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<Rc<RefCell<Vec<Value>>>>,
pub global_purity: &'a Option<Rc<RefCell<HashMap<GlobalIdx, Purity>>>>,
}
impl<'a> Inliner<'a> {
pub fn new(
globals: &'a Option<Rc<RefCell<Vec<Value>>>>,
global_purity: &'a Option<Rc<RefCell<HashMap<GlobalIdx, Purity>>>>,
) -> 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>() {
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<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);
}
}
_ => {}
}
}
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);
}
}
_ => {}
}
}
}
src/ast/compiler/optimizer/mod.rs
+3
View File
@@ -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};
src/ast/compiler/optimizer/utils.rs
+18
View File
@@ -0,0 +1,18 @@
use crate::ast::compiler::bound_nodes::{AnalyzedNode, BoundKind};
pub fn flatten_tuple(node: AnalyzedNode, into: &mut Vec<AnalyzedNode>) {
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),
}
}