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

Refactor stack size calculation

Browse Source 
Introduces a `StackAllocator` to manage local slot mapping during
lowering. This allows for efficient calculation of stack sizes for
lambdas and the root node. Tail position marking is now handled more
accurately.
This commit is contained in:
Michael Schimmel
2026-03-13 16:50:57 +01:00
parent 8f3e9cbd72
commit 222fe64a1f
1 changed files with 84 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ast/compiler/lowering.rs
+84 -117
View File
@@ -1,5 +1,6 @@
use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, Node};
use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, LocalSlot, Node};
use crate::ast::types::StaticType;
use std::collections::HashMap;
use std::fmt::Debug;
use std::rc::Rc;
@@ -25,94 +26,34 @@ impl Debug for RuntimeMetadata {
/// The ExecNode is the AST used by the VM. It carries TCO flags and links to metrics.
pub type ExecNode = Node<RuntimeMetadata>;
fn calc_stack_size<T>(root_node: &Node<T>) -> u32 {
let mut max_slot = -1i32;
struct StackAllocator {
mapping: HashMap<u32, u32>,
next_slot: u32,
}
fn visit<T>(node: &Node<T>, max_slot: &mut i32) {
match &node.kind {
BoundKind::Get {
addr: Address::Local(slot),
..
}
| BoundKind::Set {
addr: Address::Local(slot),
..
}
| BoundKind::Define {
addr: Address::Local(slot),
..
} => {
if slot.0 as i32 > *max_slot {
*max_slot = slot.0 as i32;
}
}
BoundKind::Lambda { .. } => {
// Do NOT recurse into nested lambdas to prevent over-allocating outer stacks
}
_ => {}
}
// Generic traversal
match &node.kind {
BoundKind::If {
cond,
then_br,
else_br,
} => {
visit(cond, max_slot);
visit(then_br, max_slot);
if let Some(e) = else_br {
visit(e, max_slot);
}
}
BoundKind::Set { value, .. } | BoundKind::Define { value, .. } => {
visit(value, max_slot);
}
BoundKind::Destructure { pattern, value } => {
visit(pattern, max_slot);
visit(value, max_slot);
}
BoundKind::Pipe { inputs, lambda, .. } => {
for i in inputs {
visit(i, max_slot);
}
visit(lambda, max_slot);
}
BoundKind::Call { callee, args } => {
visit(callee, max_slot);
visit(args, max_slot);
}
BoundKind::Again { args } => visit(args, max_slot),
BoundKind::Block { exprs } => {
for e in exprs {
visit(e, max_slot);
}
}
BoundKind::Tuple { elements } => {
for e in elements {
visit(e, max_slot);
}
}
BoundKind::Record { values, .. } => {
for v in values {
visit(v, max_slot);
}
}
BoundKind::Expansion { bound_expanded, .. } => visit(bound_expanded, max_slot),
_ => {}
impl StackAllocator {
fn new() -> Self {
Self {
mapping: HashMap::new(),
next_slot: 0,
}
}
// Special case for root: if the node itself is a lambda, we DO want to visit its params and body,
// but not any deeply nested lambdas.
if let BoundKind::Lambda { params, body, .. } = &root_node.kind {
visit(params, &mut max_slot);
visit(body, &mut max_slot);
} else {
visit(root_node, &mut max_slot);
fn map_slot(&mut self, slot: LocalSlot) -> LocalSlot {
let entry = self.mapping.entry(slot.0).or_insert_with(|| {
let s = self.next_slot;
self.next_slot += 1;
s
});
LocalSlot(*entry)
}
(max_slot + 1) as u32
fn map_address(&mut self, addr: Address) -> Address {
match addr {
Address::Local(slot) => Address::Local(self.map_slot(slot)),
other => other,
}
}
}
pub struct Lowering;
@@ -120,18 +61,30 @@ pub struct Lowering;
impl Lowering {
/// Lowers an AnalyzedNode to an ExecNode, marking tail positions and calculating stack sizes.
pub fn lower(node: AnalyzedNode) -> ExecNode {
let root_stack_size = calc_stack_size(&node);
let mut exec_node = Self::transform(Rc::new(node), true);
exec_node.ty.stack_size = root_stack_size;
let mut allocator = StackAllocator::new();
let mut exec_node = Self::transform(Rc::new(node), true, &mut allocator);
// If the top-level node is a Lambda, it already has its internal stack_size
// calculated during transform(). For non-lambdas (like raw expressions),
// we use the allocator's next_slot to determine the required root stack size.
if !matches!(exec_node.kind, BoundKind::Lambda { .. }) {
exec_node.ty.stack_size = allocator.next_slot;
}
exec_node
}
fn transform(node_rc: Rc<AnalyzedNode>, is_tail_position: bool) -> ExecNode {
fn transform(
node_rc: Rc<AnalyzedNode>,
is_tail_position: bool,
allocator: &mut StackAllocator,
) -> ExecNode {
let node = &*node_rc;
let mut lambda_stack_size = 0;
let new_kind = match &node.kind {
BoundKind::Call { callee, args } => BoundKind::Call {
callee: Rc::new(Self::transform(callee.clone(), false)),
args: Rc::new(Self::transform(args.clone(), false)),
callee: Rc::new(Self::transform(callee.clone(), false, allocator)),
args: Rc::new(Self::transform(args.clone(), false, allocator)),
},
BoundKind::Again { args } => {
@@ -139,7 +92,7 @@ impl Lowering {
panic!("'again' is only allowed in tail position to avoid dead code.");
}
BoundKind::Again {
args: Rc::new(Self::transform(args.clone(), false)),
args: Rc::new(Self::transform(args.clone(), false, allocator)),
}
}
BoundKind::Pipe {
@@ -149,11 +102,11 @@ impl Lowering {
} => {
let mut t_inputs = Vec::with_capacity(inputs.len());
for input in inputs {
t_inputs.push(Rc::new(Self::transform(input.clone(), false)));
t_inputs.push(Rc::new(Self::transform(input.clone(), false, allocator)));
}
BoundKind::Pipe {
inputs: t_inputs,
lambda: Rc::new(Self::transform(lambda.clone(), false)),
lambda: Rc::new(Self::transform(lambda.clone(), false, allocator)),
out_type: out_type.clone(),
}
}
@@ -163,14 +116,11 @@ impl Lowering {
then_br,
else_br,
} => BoundKind::If {
cond: Rc::new(Self::transform(cond.clone(), false)),
then_br: Rc::new(Self::transform(
then_br.clone(),
is_tail_position,
)),
cond: Rc::new(Self::transform(cond.clone(), false, allocator)),
then_br: Rc::new(Self::transform(then_br.clone(), is_tail_position, allocator)),
else_br: else_br
.as_ref()
.map(|e| Rc::new(Self::transform(e.clone(), is_tail_position))),
.map(|e| Rc::new(Self::transform(e.clone(), is_tail_position, allocator))),
},
BoundKind::Block { exprs } => {
@@ -185,6 +135,7 @@ impl Lowering {
new_exprs.push(Rc::new(Self::transform(
expr.clone(),
is_tail_position && is_last,
allocator,
)));
}
BoundKind::Block { exprs: new_exprs }
@@ -196,16 +147,30 @@ impl Lowering {
upvalues,
body,
positional_count,
} => BoundKind::Lambda {
params: Rc::new(Self::transform(params.clone(), false)),
upvalues: upvalues.clone(),
body: Rc::new(Self::transform(body.clone(), true)),
positional_count: *positional_count,
},
} => {
// Upvalues refer to the PARENT scope's addresses.
let mapped_upvalues = upvalues
.iter()
.map(|a| allocator.map_address(*a))
.collect();
// New allocator for the lambda's own stack frame
let mut lambda_allocator = StackAllocator::new();
let t_params = Rc::new(Self::transform(params.clone(), false, &mut lambda_allocator));
let t_body = Rc::new(Self::transform(body.clone(), true, &mut lambda_allocator));
lambda_stack_size = lambda_allocator.next_slot;
BoundKind::Lambda {
params: t_params,
upvalues: mapped_upvalues,
body: t_body,
positional_count: *positional_count,
}
}
BoundKind::Set { addr, value } => BoundKind::Set {
addr: *addr,
value: Rc::new(Self::transform(value.clone(), false)),
addr: allocator.map_address(*addr),
value: Rc::new(Self::transform(value.clone(), false, allocator)),
},
BoundKind::Define {
name,
@@ -215,19 +180,19 @@ impl Lowering {
captured_by,
} => BoundKind::Define {
name: name.clone(),
addr: *addr,
addr: allocator.map_address(*addr),
kind: *kind,
value: Rc::new(Self::transform(value.clone(), false)),
value: Rc::new(Self::transform(value.clone(), false, allocator)),
captured_by: captured_by.clone(),
},
BoundKind::Destructure { pattern, value } => BoundKind::Destructure {
pattern: Rc::new(Self::transform(pattern.clone(), false)),
value: Rc::new(Self::transform(value.clone(), false)),
pattern: Rc::new(Self::transform(pattern.clone(), false, allocator)),
value: Rc::new(Self::transform(value.clone(), false, allocator)),
},
BoundKind::Record { layout, values } => {
let new_values = values
.iter()
.map(|v| Rc::new(Self::transform(v.clone(), false)))
.map(|v| Rc::new(Self::transform(v.clone(), false, allocator)))
.collect();
BoundKind::Record {
layout: layout.clone(),
@@ -237,7 +202,7 @@ impl Lowering {
BoundKind::Tuple { elements } => {
let new_elements = elements
.iter()
.map(|e| Rc::new(Self::transform(e.clone(), false)))
.map(|e| Rc::new(Self::transform(e.clone(), false, allocator)))
.collect();
BoundKind::Tuple {
elements: new_elements,
@@ -245,12 +210,12 @@ impl Lowering {
}
BoundKind::Constant(v) => BoundKind::Constant(v.clone()),
BoundKind::Get { addr, name } => BoundKind::Get {
addr: *addr,
addr: allocator.map_address(*addr),
name: name.clone(),
},
BoundKind::FieldAccessor(k) => BoundKind::FieldAccessor(*k),
BoundKind::GetField { rec, field } => BoundKind::GetField {
rec: Rc::new(Self::transform(rec.clone(), false)),
rec: Rc::new(Self::transform(rec.clone(), false, allocator)),
field: *field,
},
BoundKind::Nop => BoundKind::Nop,
@@ -262,15 +227,17 @@ impl Lowering {
bound_expanded: Rc::new(Self::transform(
bound_expanded.clone(),
is_tail_position,
allocator,
)),
},
BoundKind::Extension(_) => BoundKind::Nop,
BoundKind::Error => BoundKind::Error,
};
let stack_size = match &new_kind {
BoundKind::Lambda { .. } => calc_stack_size(node),
_ => 0,
let stack_size = if let BoundKind::Lambda { .. } = &new_kind {
lambda_stack_size
} else {
0
};
Node {