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

@@ -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 {
+struct StackAllocator {
-    let mut max_slot = -1i32;
+    mapping: HashMap<u32, u32>,
+    next_slot: u32,
+}
 
-    fn visit<T>(node: &Node<T>, max_slot: &mut i32) {
+impl StackAllocator {
-        match &node.kind {
+    fn new() -> Self {
-            BoundKind::Get {
+        Self {
-                addr: Address::Local(slot),
+            mapping: HashMap::new(),
-                ..
+            next_slot: 0,
-            }
-            | 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,
+    fn map_slot(&mut self, slot: LocalSlot) -> LocalSlot {
-    // but not any deeply nested lambdas.
+        let entry = self.mapping.entry(slot.0).or_insert_with(|| {
-    if let BoundKind::Lambda { params, body, .. } = &root_node.kind {
+            let s = self.next_slot;
-        visit(params, &mut max_slot);
+            self.next_slot += 1;
-        visit(body, &mut max_slot);
+            s
-    } else {
+        });
-        visit(root_node, &mut max_slot);
+        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 allocator = StackAllocator::new();
-        let mut exec_node = Self::transform(Rc::new(node), true);
+        let mut exec_node = Self::transform(Rc::new(node), true, &mut allocator);
-        exec_node.ty.stack_size = root_stack_size;
+
+        // 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)),
+                callee: Rc::new(Self::transform(callee.clone(), false, allocator)),
-                args: Rc::new(Self::transform(args.clone(), false)),
+                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)),
+                cond: Rc::new(Self::transform(cond.clone(), false, allocator)),
-                then_br: Rc::new(Self::transform(
+                then_br: Rc::new(Self::transform(then_br.clone(), is_tail_position, allocator)),
-                    then_br.clone(),
-                    is_tail_position,
-                )),
                 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 refer to the PARENT scope's addresses.
-                upvalues: upvalues.clone(),
+                let mapped_upvalues = upvalues
-                body: Rc::new(Self::transform(body.clone(), true)),
+                    .iter()
-                positional_count: *positional_count,
+                    .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,
+                addr: allocator.map_address(*addr),
-                value: Rc::new(Self::transform(value.clone(), false)),
+                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)),
+                pattern: Rc::new(Self::transform(pattern.clone(), false, allocator)),
-                value: Rc::new(Self::transform(value.clone(), false)),
+                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 {
+        let stack_size = if let BoundKind::Lambda { .. } = &new_kind {
-            BoundKind::Lambda { .. } => calc_stack_size(node),
+            lambda_stack_size
-            _ => 0,
+        } else {
+            0
         };
 
         Node {