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

Add example and refine local inlining

Browse Source 
Add a new example file `examples/def_local_inlining.myc` to demonstrate
and test local inlining.

Refine the optimizer to handle local inlining more robustly by:
- Passing a `SubstitutionMap` to `visit_node` to track local variable
  substitutions.
- Enabling constant propagation for local variables by checking
  `sub.locals` in `BoundKind::Get`.
- Updating `BoundKind::DefLocal` and `BoundKind::Set` to maintain the
  substitution map for local variables.
- Adjusting `try_beta_reduce` to use the optimizer's `visit_node` with
  the substitution map for inlining.
- Re-indexing upvalues correctly when inlining captured variables into
  lambdas.
This commit is contained in:
Michael Schimmel
2026-02-21 22:35:04 +01:00
parent ff1024ee49
commit b2e010e755
4 changed files with 169 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/def_local_inlining.myc
+23
View File
@@ -0,0 +1,23 @@
;; examples/def_local_inlining.myc
;; Demonstrates potential for DefLocal-Inlining (Phase 2.5)
(fn []
(do
;; 1. Simple constant propagation
(def x 10)
(def y 20)
;; Expected optimization: (+ 10 20) -> 30
(def z (+ x y))
;; 2. Tracking assignments (Assign-Safety)
;; Inlining must be careful with 'assign'.
(def result
(do
(def a 100)
(assign a 200)
;; Expected optimization: 200
a))
;; 3. Return a tuple of optimized results
[z result]))
src/ast/compiler/optimizer.rs
+130 -162
View File
@@ -23,7 +23,8 @@ impl Optimizer {
let mut current = node;
for _ in 0..self.max_passes {
let next = self.visit_node(current.clone());
let mut sub = SubstitutionMap::new();
let next = self.visit_node(current.clone(), &mut sub);
if next == current {
break;
}
@@ -32,32 +33,55 @@ impl Optimizer {
current
}
fn visit_node(&self, node: TypedNode) -> TypedNode {
fn visit_node(&self, node: TypedNode, sub: &mut SubstitutionMap) -> TypedNode {
let (new_kind, new_ty) = match node.kind {
BoundKind::Get { addr, name } => {
match addr {
Address::Local(slot) => {
if let Some(val) = sub.locals.get(&slot) {
return Node {
identity: node.identity,
ty: val.static_type(),
kind: BoundKind::Constant(val.clone()),
};
}
}
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()),
};
}
}
_ => {}
}
(BoundKind::Get { addr, name }, node.ty)
}
BoundKind::Call { callee, args } => {
let callee = self.visit_node(*callee);
let args = self.visit_node(*args);
let callee = self.visit_node(*callee, sub);
let args = self.visit_node(*args, sub);
if self.level >= 2 {
// Case 1: Beta-Reduction for Lambda Literals
if let BoundKind::Lambda { params, body, .. } = &callee.kind
&& let Some(collapsed) = self.try_beta_reduce(params, &args, (**body).clone()) {
return self.visit_node(collapsed);
return self.visit_node(collapsed, sub);
}
// Case 2: Cracking and Inlining for Constant Closures
if let BoundKind::Constant(Value::Object(ref obj)) = callee.kind
&& let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
let mut sub = SubstitutionMap::new();
let mut closure_sub = SubstitutionMap::new();
for (i, cell) in closure.upvalues.iter().enumerate() {
sub.add_upvalue(i as u32, cell.borrow().clone());
closure_sub.add_upvalue(i as u32, cell.borrow().clone());
}
// DIRECT ACCESS: Use the original TypedNode from the closure
let inlined_body = sub.inline((*closure.function_node).clone());
let inlined_body = self.visit_node((*closure.function_node).clone(), &mut closure_sub);
// Try to collapse the now-stateless lambda call
if let Some(collapsed) = self.try_beta_reduce(&closure.parameter_node, &args, inlined_body) {
return self.visit_node(collapsed);
return self.visit_node(collapsed, sub);
}
}
@@ -66,23 +90,21 @@ impl Optimizer {
}
}
// Level 1: Just Crack closures but keep the Call structure
if self.level >= 1
&& let BoundKind::Constant(Value::Object(ref obj)) = callee.kind
&& let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
let mut sub = SubstitutionMap::new();
let mut closure_sub = SubstitutionMap::new();
for (i, cell) in closure.upvalues.iter().enumerate() {
sub.add_upvalue(i as u32, cell.borrow().clone());
closure_sub.add_upvalue(i as u32, cell.borrow().clone());
}
// DIRECT ACCESS: Use the original TypedNode
let inlined_body = sub.inline((*closure.function_node).clone());
let inlined_body = self.visit_node((*closure.function_node).clone(), &mut closure_sub);
let cracked_lambda = Node {
identity: callee.identity.clone(),
ty: callee.ty.clone(),
kind: BoundKind::Lambda {
params: closure.parameter_node.clone(),
upvalues: vec![], // CRACKED: Stateless
upvalues: vec![],
body: Rc::new(inlined_body),
positional_count: closure.positional_count,
},
@@ -98,27 +120,27 @@ impl Optimizer {
},
BoundKind::If { cond, then_br, else_br } => {
let cond = self.visit_node(*cond);
let cond = self.visit_node(*cond, sub);
if self.level >= 2
&& let BoundKind::Constant(ref val) = cond.kind {
if val.is_truthy() {
return self.visit_node(*then_br);
return self.visit_node(*then_br, sub);
} else if let Some(else_node) = else_br {
return self.visit_node(*else_node);
return self.visit_node(*else_node, sub);
} else {
return Node { identity: node.identity, kind: BoundKind::Nop, ty: StaticType::Void };
}
}
let then_br = Box::new(self.visit_node(*then_br));
let else_br = else_br.map(|e| Box::new(self.visit_node(*e)));
let then_br = Box::new(self.visit_node(*then_br, sub));
let else_br = else_br.map(|e| Box::new(self.visit_node(*e, sub)));
(BoundKind::If { cond: Box::new(cond), then_br, else_br }, node.ty)
},
BoundKind::Block { exprs } => {
let mut new_exprs = Vec::with_capacity(exprs.len());
for e in exprs {
let opt = self.visit_node(e);
let opt = self.visit_node(e, sub);
if self.level >= 2 && matches!(opt.kind, BoundKind::Nop) {
continue;
}
@@ -137,35 +159,83 @@ impl Optimizer {
(BoundKind::Block { exprs: new_exprs }, ty)
},
BoundKind::Lambda { params, upvalues, body, positional_count } => {
let params = Rc::new(self.visit_node(params.as_ref().clone()));
let body = Rc::new(self.visit_node(body.as_ref().clone()));
(BoundKind::Lambda { params, upvalues, body, positional_count }, node.ty)
BoundKind::Lambda { params, upvalues: original_upvalues, body, positional_count } => {
let mut new_upvalues = Vec::new();
let mut mapping = Vec::new();
let mut next_inner_subs = SubstitutionMap::new();
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()); }
}
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));
new_upvalues.push(*capture_addr);
}
}
let params = Rc::new(self.visit_node(params.as_ref().clone(), &mut SubstitutionMap::new()));
let body_node = self.visit_node((*body).clone(), &mut next_inner_subs);
let reindexed_body = if new_upvalues.len() != original_upvalues.len() {
sub.reindex_upvalues(body_node, &mapping)
} else {
body_node
};
(BoundKind::Lambda {
params,
upvalues: new_upvalues,
body: Rc::new(reindexed_body),
positional_count
}, node.ty)
},
// Standard Recursive Traversal
BoundKind::DefLocal { name, slot, value, captured_by } => {
let value = Box::new(self.visit_node(*value));
let value = Box::new(self.visit_node(*value, sub));
if let BoundKind::Constant(val) = &value.kind {
sub.add_local(slot, val.clone());
} else {
sub.locals.remove(&slot);
}
(BoundKind::DefLocal { name, slot, value, captured_by }, node.ty)
},
BoundKind::DefGlobal { name, global_index, value } => {
let value = Box::new(self.visit_node(*value));
let value = Box::new(self.visit_node(*value, sub));
(BoundKind::DefGlobal { name, global_index, value }, node.ty)
},
BoundKind::Set { addr, value } => {
let value = Box::new(self.visit_node(*value));
let value = Box::new(self.visit_node(*value, sub));
if let Address::Local(slot) = addr {
if let BoundKind::Constant(val) = &value.kind {
sub.add_local(slot, val.clone());
} else {
sub.locals.remove(&slot);
}
}
(BoundKind::Set { addr, value }, node.ty)
},
BoundKind::Tuple { elements } => {
let elements = elements.into_iter().map(|e| self.visit_node(e)).collect();
let elements = elements.into_iter().map(|e| self.visit_node(e, sub)).collect();
(BoundKind::Tuple { elements }, node.ty)
},
BoundKind::Record { fields } => {
let fields = fields.into_iter().map(|(k, v)| (self.visit_node(k), self.visit_node(v))).collect();
let fields = fields.into_iter().map(|(k, v)| (self.visit_node(k, sub), self.visit_node(v, sub))).collect();
(BoundKind::Record { fields }, node.ty)
},
BoundKind::Expansion { original_call, bound_expanded } => {
let bound_expanded = Box::new(self.visit_node(*bound_expanded));
let bound_expanded = Box::new(self.visit_node(*bound_expanded, sub));
(BoundKind::Expansion { original_call, bound_expanded }, node.ty)
},
k => (k, node.ty),
@@ -175,7 +245,6 @@ impl Optimizer {
}
fn try_beta_reduce(&self, params: &TypedNode, args: &TypedNode, body: TypedNode) -> Option<TypedNode> {
// SAFETY: Only beta-reduce if the body does not contain local definitions.
if self.contains_def_local(&body) {
return None;
}
@@ -187,17 +256,15 @@ impl Optimizer {
let mut slot_index = 0;
self.map_params_to_args(params, &arg_vals, &mut slot_index, &mut sub);
// SAFETY: Only beta-reduce if ALL provided arguments were successfully inlined as constants.
// If we leave any Local(i) references, the VM will crash because the frame is gone.
if sub.locals.len() < arg_vals.len() {
return None;
}
if sub.is_empty() && !arg_vals.is_empty() {
if sub.locals.is_empty() && sub.upvalues.is_empty() && !arg_vals.is_empty() {
return None;
}
Some(sub.inline(body))
Some(self.visit_node(body, &mut sub))
}
fn flatten_typed_tuple(&self, node: &TypedNode, into: &mut Vec<TypedNode>) {
@@ -272,7 +339,7 @@ impl Optimizer {
}
BoundKind::Block { exprs } => exprs.iter().any(|e| self.contains_def_local(e)),
BoundKind::Call { callee, args } => self.contains_def_local(callee) || self.contains_def_local(args),
BoundKind::Lambda { .. } => false, // Nested lambda definitions are safe (they have their own frame)
BoundKind::Lambda { .. } => false,
BoundKind::Tuple { elements } => elements.iter().any(|e| self.contains_def_local(e)),
BoundKind::Record { fields } => fields.iter().any(|(k, v)| self.contains_def_local(k) || self.contains_def_local(v)),
BoundKind::Set { value, .. } => self.contains_def_local(value),
@@ -302,141 +369,42 @@ impl SubstitutionMap {
self.upvalues.insert(idx, val);
}
fn is_empty(&self) -> bool {
self.locals.is_empty() && self.upvalues.is_empty()
}
fn inline(&self, node: TypedNode) -> TypedNode {
self.inline_recursive(node, 0, &HashMap::new())
}
fn inline_recursive(&self, node: TypedNode, depth: usize, upvalue_subs: &HashMap<u32, Value>) -> TypedNode {
let (new_kind, new_ty) = match node.kind {
BoundKind::Get { addr, name } => {
match addr {
Address::Local(slot) if depth == 0 => {
if let Some(val) = self.locals.get(&slot) {
(BoundKind::Constant(val.clone()), val.static_type())
} else { (BoundKind::Get { addr, name }, node.ty) }
}
Address::Upvalue(idx) if depth == 0 => {
if let Some(val) = self.upvalues.get(&idx) {
(BoundKind::Constant(val.clone()), val.static_type())
} else { (BoundKind::Get { addr, name }, node.ty) }
}
Address::Upvalue(idx) if depth > 0 => {
if let Some(val) = upvalue_subs.get(&idx) {
(BoundKind::Constant(val.clone()), val.static_type())
} else { (BoundKind::Get { addr, name }, node.ty) }
}
_ => (BoundKind::Get { addr, name }, node.ty)
}
},
BoundKind::Lambda { params, upvalues: original_upvalues, body, positional_count } => {
// Phase 3: Un-Capturing
// If this lambda captures variables that we just inlined, we MUST remove them
// from its capture list to avoid "Stack underflow capture" errors in the VM.
let mut new_upvalues = Vec::new();
let mut next_inner_subs = HashMap::new();
let mut mapping = Vec::new(); // Old Upvalue Index -> New Upvalue Index (or None)
for (old_idx, capture_addr) in original_upvalues.iter().enumerate() {
let mut inlined_val = None;
if depth == 0 {
match capture_addr {
Address::Local(slot) => {
if let Some(val) = self.locals.get(slot) { inlined_val = Some(val.clone()); }
}
Address::Upvalue(idx) => {
if let Some(val) = self.upvalues.get(idx) { inlined_val = Some(val.clone()); }
}
_ => {}
}
} else if let Address::Upvalue(old_parent_idx) = capture_addr
&& let Some(val) = upvalue_subs.get(old_parent_idx) {
inlined_val = Some(val.clone());
}
if let Some(val) = inlined_val {
next_inner_subs.insert(old_idx as u32, val);
mapping.push(None);
} else {
mapping.push(Some(new_upvalues.len() as u32));
new_upvalues.push(*capture_addr);
}
}
// If any upvalues were removed, we must re-index the Get(Upvalue) nodes in the body.
let body_node = self.inline_recursive(body.as_ref().clone(), depth + 1, &next_inner_subs);
let reindexed_body = if new_upvalues.len() != original_upvalues.len() {
self.reindex_upvalues(body_node, &mapping)
} else {
body_node
};
(BoundKind::Lambda { params, upvalues: new_upvalues, body: Rc::new(reindexed_body), positional_count }, node.ty)
},
// Boilerplate Recursive Traversal
BoundKind::If { cond, then_br, else_br } => {
let cond = Box::new(self.inline_recursive(*cond, depth, upvalue_subs));
let then_br = Box::new(self.inline_recursive(*then_br, depth, upvalue_subs));
let else_br = else_br.map(|e| Box::new(self.inline_recursive(*e, depth, upvalue_subs)));
(BoundKind::If { cond, then_br, else_br }, node.ty)
},
BoundKind::Block { exprs } => {
let exprs = exprs.into_iter().map(|e| self.inline_recursive(e, depth, upvalue_subs)).collect();
(BoundKind::Block { exprs }, node.ty)
},
BoundKind::Call { callee, args } => {
let callee = Box::new(self.inline_recursive(*callee, depth, upvalue_subs));
let args = Box::new(self.inline_recursive(*args, depth, upvalue_subs));
(BoundKind::Call { callee, args }, node.ty)
},
BoundKind::DefLocal { name, slot, value, captured_by } => {
let value = Box::new(self.inline_recursive(*value, depth, upvalue_subs));
(BoundKind::DefLocal { name, slot, value, captured_by }, node.ty)
},
BoundKind::Set { addr, value } => {
let value = Box::new(self.inline_recursive(*value, depth, upvalue_subs));
(BoundKind::Set { addr, value }, node.ty)
},
BoundKind::Tuple { elements } => {
let elements = elements.into_iter().map(|e| self.inline_recursive(e, depth, upvalue_subs)).collect();
(BoundKind::Tuple { elements }, node.ty)
},
BoundKind::Record { fields } => {
let fields = fields.into_iter().map(|(k, v)| (self.inline_recursive(k, depth, upvalue_subs), self.inline_recursive(v, depth, upvalue_subs))).collect();
(BoundKind::Record { fields }, node.ty)
},
k => (k, node.ty),
};
Node { identity: node.identity, kind: new_kind, ty: new_ty }
}
/// Re-maps Get(Upvalue(old_idx)) to Get(Upvalue(new_idx)) based on a mapping table.
fn reindex_upvalues(&self, node: TypedNode, mapping: &[Option<u32>]) -> TypedNode {
let (new_kind, new_ty) = match node.kind {
BoundKind::Get { addr: Address::Upvalue(idx), name } => {
if let Some(Some(new_idx)) = mapping.get(idx as usize) {
(BoundKind::Get { addr: Address::Upvalue(*new_idx), name }, node.ty)
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), // Should have been inlined
}
} else {
(BoundKind::Get { addr: Address::Upvalue(idx), name }, node.ty)
}
},
// Nested Lambda: Increment depth (we only re-index upvalues of the current scope)
BoundKind::Lambda { params, upvalues, body, positional_count } => {
// IMPORTANT: In MyC, nested upvalues refer to the capture list of the nested lambda.
// We DON'T re-index inside nested lambdas because their upvalue list is separate.
(BoundKind::Lambda { params, upvalues, body, positional_count }, node.ty)
// IMPORTANT: If this nested lambda captures an upvalue from our current scope,
// we MUST re-index it in its own capture list!
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; // Inlined or re-indexed
}
next_upvalues.push(addr);
}
// Note: We don't recurse into the body with the SAME mapping,
// because nested Get(Upvalue) nodes refer to THIS lambda's capture list.
(BoundKind::Lambda { params, upvalues: next_upvalues, body, positional_count }, node.ty)
},
// Recursive Traversal
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)));
src/ast/compiler/tco.rs
+12 -1
View File
@@ -3,7 +3,9 @@ use crate::ast::nodes::Node;
use crate::ast::compiler::bound_nodes::{BoundKind, TypedNode};
use crate::ast::types::StaticType;
#[derive(Debug, Clone)]
use std::fmt::Debug;
#[derive(Clone)]
pub struct RuntimeMetadata {
pub ty: StaticType,
pub is_tail: bool,
@@ -11,6 +13,15 @@ pub struct RuntimeMetadata {
pub original: Rc<TypedNode>,
}
impl Debug for RuntimeMetadata {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Metadata")
.field("ty", &self.ty)
.field("is_tail", &self.is_tail)
.finish()
}
}
/// The ExecNode is the AST used by the VM. It carries TCO flags and links to source.
pub type ExecNode = Node<BoundKind<RuntimeMetadata>, RuntimeMetadata>;
src/utils/tester.rs
+4 -1
View File
@@ -52,7 +52,10 @@ pub fn run_functional_tests_with_level(level: u32) -> Vec<TestResult> {
results.push(TestResult {
name,
success: false,
message: format!("Level {}: Expected {}, got {}", level, expected, val_str),
message: format!(
"Level {}: Expected {}, got {}",
level, expected, val_str
),
});
}
}