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

Extract calculate_stack_size to a standalone function

Browse Source 
The `calculate_stack_size` method was duplicated in `bound_nodes.rs` and
`tco.rs`. This commit extracts it into a single standalone function in
`tco.rs` to avoid duplication. The stack size is now calculated and
stored in the `ExecNode`'s `ty` field during the TCO optimization phase.
This commit is contained in:
Michael Schimmel
2026-03-11 10:35:18 +01:00
parent 3657f19047
commit db26719cad
4 changed files with 106 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ast/compiler/bound_nodes.rs
-93
View File
@@ -304,99 +304,6 @@ where
/// A single field in a Record literal (Key-Value pair)
pub type RecordField<T> = (BoundNode<T>, BoundNode<T>);
impl<T> BoundNode<T> {
pub fn calculate_stack_size(&self) -> u32 {
let mut max_slot = -1i32;
fn visit<T>(node: &BoundNode<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;
}
}
_ => {}
}
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, .. } => {
visit(value, max_slot);
}
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),
BoundKind::Lambda { params, body, .. } => {
// Check parameters and body of the lambda,
// but do NOT recurse into nested lambdas (which is handled by the generic recursion blocker below).
visit(params, max_slot);
visit(body, max_slot);
}
_ => {}
}
}
// Handle the root node: if it's a lambda, we want to check its content.
// If we just called visit(self), it would hit the Lambda case.
visit(self, &mut max_slot);
(max_slot + 1) as u32
}
}
impl<T> BoundKind<T> {
pub fn display_name(&self) -> String {
match self {
src/ast/compiler/tco.rs
+96 -8
View File
@@ -1,4 +1,4 @@
use crate::ast::compiler::bound_nodes::{AnalyzedNode, BoundKind};
use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind};
use crate::ast::nodes::Node;
use crate::ast::types::StaticType;
use std::fmt::Debug;
@@ -26,12 +26,105 @@ 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<BoundKind<RuntimeMetadata>, RuntimeMetadata>;
fn calc_stack_size<T>(root_node: &Node<BoundKind<T>, T>) -> u32 {
let mut max_slot = -1i32;
fn visit<T>(node: &Node<BoundKind<T>, 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),
_ => {}
}
}
// 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);
}
(max_slot + 1) as u32
}
pub struct TCO;
impl TCO {
/// Lowers an AnalyzedNode to an ExecNode and marks tail positions.
pub fn optimize(node: AnalyzedNode) -> ExecNode {
Self::transform(Rc::new(node), true)
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;
exec_node
}
fn transform(node_rc: Rc<AnalyzedNode>, is_tail_position: bool) -> ExecNode {
@@ -177,12 +270,7 @@ impl TCO {
};
let stack_size = match &new_kind {
BoundKind::Lambda { .. } => {
// Pre-calculate stack size for the lambda body.
// Note: 'node' here is AnalyzedNode (Node<BoundKind<NodeMetrics>, NodeMetrics>)
// 'calculate_stack_size' works for any T.
node.calculate_stack_size()
}
BoundKind::Lambda { .. } => calc_stack_size(node),
_ => 0,
};