diff --git a/src/ast/compiler/lowering.rs b/src/ast/compiler/lowering.rs index cf8ecc4..c20bc63 100644 --- a/src/ast/compiler/lowering.rs +++ b/src/ast/compiler/lowering.rs @@ -8,6 +8,8 @@ use std::rc::Rc; struct StackAllocator { mapping: HashMap, + /// Physical slots that have been freed and can be reused. + free_slots: Vec, next_slot: u32, } @@ -15,15 +17,21 @@ impl StackAllocator { fn new() -> Self { Self { mapping: HashMap::new(), + free_slots: Vec::new(), next_slot: 0, } } fn map_slot(&mut self, slot: VirtualId) -> StackOffset { let entry = self.mapping.entry(slot.0).or_insert_with(|| { - let s = self.next_slot; - self.next_slot += 1; - s + // Reuse a freed slot before growing the frame. + if let Some(recycled) = self.free_slots.pop() { + recycled + } else { + let s = self.next_slot; + self.next_slot += 1; + s + } }); StackOffset(*entry) } @@ -35,6 +43,15 @@ impl StackAllocator { Address::Global(idx) => Address::Global(idx), } } + + /// Releases the physical slot of `vid` back to the free-list. + /// If `vid` was never mapped (e.g. a dead def removed by the optimizer) + /// this is a no-op. + fn free_slot(&mut self, vid: VirtualId) { + if let Some(&physical) = self.mapping.get(&vid.0) { + self.free_slots.push(physical); + } + } } pub struct Lowering; @@ -116,6 +133,12 @@ impl Lowering { if exprs.is_empty() { NodeKind::Block { exprs: vec![] } } else { + // Collect VirtualIds of non-captured locals before + // transforming. After all exprs are lowered their + // physical slots can be returned to the free-list, + // allowing subsequent blocks to reuse them. + let scope_locals = Self::collect_scope_locals(exprs); + let last_idx = exprs.len() - 1; let mut new_exprs = Vec::with_capacity(exprs.len()); @@ -127,6 +150,13 @@ impl Lowering { allocator, ))); } + + // Release slots — variables defined here are no longer + // live after this block ends. + for vid in scope_locals { + allocator.free_slot(vid); + } + NodeKind::Block { exprs: new_exprs } } } @@ -268,4 +298,37 @@ impl Lowering { comments: node.comments.clone(), } } + + /// Collects the `VirtualId`s of all non-captured locals defined at the + /// top level of `exprs`. These slots are safe to reclaim once the block + /// ends. Captured locals (upvalues) are excluded — closures that outlive + /// the block still hold live references to those physical slots. + fn collect_scope_locals(exprs: &[Rc]) -> Vec { + let mut locals = Vec::new(); + for expr in exprs { + if let NodeKind::Def { pattern, info, .. } = &expr.kind + && info.captured_by.is_empty() + { + Self::collect_pattern_vids(pattern, &mut locals); + } + } + locals + } + + /// Recursively extracts `VirtualId`s from declaration identifiers in a + /// binding pattern. Handles flat (`x`) and destructuring (`[x y]`) forms. + fn collect_pattern_vids(pattern: &AnalyzedNode, out: &mut Vec) { + match &pattern.kind { + NodeKind::Identifier { + binding: IdentifierBinding::Declaration { addr: Address::Local(vid), .. }, + .. + } => out.push(*vid), + NodeKind::Tuple { elements } => { + for e in elements { + Self::collect_pattern_vids(e, out); + } + } + _ => {} + } + } } diff --git a/tests/optimizer.rs b/tests/optimizer.rs index 3b49cca..346e4d0 100644 --- a/tests/optimizer.rs +++ b/tests/optimizer.rs @@ -180,3 +180,39 @@ fn test_reproduce_inlining_slot_clash_crash() { let res = env.run_script(source); assert!(res.is_ok(), "Inlining slot clash triggered: {:?}", res.err()); } + +// ── Slot reuse ──────────────────────────────────────────────────────────────── + +#[test] +fn test_slot_reuse_non_overlapping_scopes() { + let mut env = Environment::new(); + env.optimization = true; + + // Two sequential inner blocks each bind a mutable local. + // assign prevents the optimizer from inlining the variables away, + // so both slots survive into the lowering pass. + // x lives only in the first block; once that block ends its slot is free. + // y is defined in the second block and must reuse x's slot. + // Expected result: the function returns 3 (y = 2+1). + let source = r#" + (fn [] + (do + (do (def x 1) (assign x (+ x 1)) x) + (do (def y 2) (assign y (+ y 1)) y))) + "#; + + // Correctness: calling the function must still produce the right value. + let result = env + .run_script("((fn [] (do (do (def x 1) (assign x (+ x 1)) x) (do (def y 2) (assign y (+ y 1)) y))))") + .unwrap(); + assert_eq!(format!("{}", result), "3"); + + // Stack efficiency: with slot reuse the Lambda's stack_size must be 1, + // not 2. The dump format is "stack_size: N" in the Lambda metadata line. + let dump = env.dump_ast(source).unwrap(); + assert!( + dump.contains("stack_size: 1"), + "Non-overlapping scopes must reuse the same slot (expected stack_size: 1). Dump:\n{}", + dump + ); +}