From 222fe64a1fe2e443ffb77ee010877333ba6ec846 Mon Sep 17 00:00:00 2001 From: Michael Schimmel Date: Fri, 13 Mar 2026 16:50:57 +0100 Subject: [PATCH] Refactor stack size calculation 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. --- src/ast/compiler/lowering.rs | 201 +++++++++++++++-------------------- 1 file changed, 84 insertions(+), 117 deletions(-) diff --git a/src/ast/compiler/lowering.rs b/src/ast/compiler/lowering.rs index e867b61..3588707 100644 --- a/src/ast/compiler/lowering.rs +++ b/src/ast/compiler/lowering.rs @@ -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; -fn calc_stack_size(root_node: &Node) -> u32 { - let mut max_slot = -1i32; +struct StackAllocator { + mapping: HashMap, + next_slot: u32, +} - fn visit(node: &Node, 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, is_tail_position: bool) -> ExecNode { + fn transform( + node_rc: Rc, + 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 {