From db26719cad85b2eba94ed37d1919d0303deb5fbf Mon Sep 17 00:00:00 2001 From: Michael Schimmel Date: Wed, 11 Mar 2026 10:35:18 +0100 Subject: [PATCH] Extract `calculate_stack_size` to a standalone function 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. --- src/ast/compiler/bound_nodes.rs | 93 ---------------------------- src/ast/compiler/tco.rs | 104 +++++++++++++++++++++++++++++--- src/ast/environment.rs | 21 +++---- src/ast/vm.rs | 7 +-- 4 files changed, 106 insertions(+), 119 deletions(-) diff --git a/src/ast/compiler/bound_nodes.rs b/src/ast/compiler/bound_nodes.rs index 3099147..acf1bb6 100644 --- a/src/ast/compiler/bound_nodes.rs +++ b/src/ast/compiler/bound_nodes.rs @@ -304,99 +304,6 @@ where /// A single field in a Record literal (Key-Value pair) pub type RecordField = (BoundNode, BoundNode); -impl BoundNode { - pub fn calculate_stack_size(&self) -> u32 { - let mut max_slot = -1i32; - - fn visit(node: &BoundNode, 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 BoundKind { pub fn display_name(&self) -> String { match self { diff --git a/src/ast/compiler/tco.rs b/src/ast/compiler/tco.rs index 71aac68..fa72a5f 100644 --- a/src/ast/compiler/tco.rs +++ b/src/ast/compiler/tco.rs @@ -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, RuntimeMetadata>; +fn calc_stack_size(root_node: &Node, T>) -> u32 { + let mut max_slot = -1i32; + + fn visit(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), + _ => {} + } + } + + // 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, 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, NodeMetrics>) - // 'calculate_stack_size' works for any T. - node.calculate_stack_size() - } + BoundKind::Lambda { .. } => calc_stack_size(node), _ => 0, }; diff --git a/src/ast/environment.rs b/src/ast/environment.rs index a6c9e73..bf519c3 100644 --- a/src/ast/environment.rs +++ b/src/ast/environment.rs @@ -144,8 +144,7 @@ impl MacroEvaluator for RuntimeMacroEvaluator { let exec_ast = TCO::optimize(Analyzer::analyze(&typed_ast, &HashMap::new())); // Minimal analysis for macro eval let mut vm = VM::new(self.global_values.clone()); - let stack_size = exec_ast.calculate_stack_size(); - vm.run(&exec_ast, stack_size) + vm.run(&exec_ast) } } @@ -601,7 +600,7 @@ impl Environment { } = &node.kind && upvalues.is_empty() { - let stack_size = node.calculate_stack_size(); + let stack_size = node.ty.stack_size; let closure = Rc::new(crate::ast::vm::Closure::new( params.ty.original.clone(), body.ty.original.clone(), @@ -629,7 +628,7 @@ impl Environment { purity: Purity::Impure, func: Rc::new(move |args| { let mut vm = VM::new(global_values.clone()); - let res = match vm.run(&exec_node, 0) { // Root call, use 0 if node is not a lambda + let res = match vm.run(&exec_node) { Ok(v) => v, Err(e) => panic!("Myc Runtime Error: {}", e), }; @@ -708,8 +707,7 @@ impl Environment { let tco_ast = TCO::optimize(optimized_ast); let mut vm = VM::new(global_values.clone()); - let stack_size = tco_ast.calculate_stack_size(); - let compiled_val = match vm.run(&tco_ast, stack_size) { + let compiled_val = match vm.run(&tco_ast) { Ok(v) => v, Err(e) => return Err(format!("VM Error during specialization: {}", e)), }; @@ -746,7 +744,6 @@ impl Environment { pub fn run_script_compiled(&self, compiled: TypedNode) -> Result { let linked = self.link(compiled); - let _stack_size = linked.calculate_stack_size(); let func = self.instantiate(linked); let res = (func.func)(&[]); self.run_pipeline(); @@ -757,16 +754,12 @@ impl Environment { self.preload_dependencies(source, None)?; let compiled = self.compile(source).into_result()?; let linked = self.link(compiled); - - // Late Counting: Determine stack size after all optimizations are finished. - // This keeps AST nodes clean and ensures the VM always has enough space. - let stack_size = linked.calculate_stack_size(); + let mut vm = VM::new(self.global_values.clone()); let mut observer = TracingObserver::new(); - - // 1. Run the script wrapper (returns a closure representing the script) - let result = vm.run_with_observer(&mut observer, &linked, stack_size); + // 1. Run the script wrapper (returns a closure representing the script) + let result = vm.run_with_observer(&mut observer, &linked); // 2. Execute the root closure immediately to get the actual script result. // All Myc scripts are wrapped in a parameterless lambda for consistency. let mut final_result = result; diff --git a/src/ast/vm.rs b/src/ast/vm.rs index 249d82d..520b23e 100644 --- a/src/ast/vm.rs +++ b/src/ast/vm.rs @@ -141,8 +141,8 @@ impl VM { } } - pub fn run(&mut self, root: &ExecNode, stack_size: u32) -> Result { - self.run_with_observer(&mut NoOpObserver, root, stack_size) + pub fn run(&mut self, root: &ExecNode) -> Result { + self.run_with_observer(&mut NoOpObserver, root) } pub fn resolve_tail_calls( @@ -263,12 +263,11 @@ impl VM { &mut self, observer: &mut O, root: &ExecNode, - stack_size: u32, ) -> Result { self.stack.clear(); self.frames.clear(); - self.stack.resize(stack_size as usize, Value::Void); + self.stack.resize(root.ty.stack_size as usize, Value::Void); self.frames.push(CallFrame { stack_base: 0,