Refactor stack allocator to reuse freed slots
Introduce `free_slots` in `StackAllocator` to keep track of physical slots that have been freed. When mapping a new virtual slot, the allocator first attempts to reuse a slot from the `free_slots` vector before allocating a new one. This change also includes a new function `collect_scope_locals` that identifies non-captured local variables defined within a block. After these expressions are lowered, the allocator reclaims the physical slots associated with these locals, making them available for reuse by subsequent blocks. A new test `test_slot_reuse_non_overlapping_scopes` is added to verify that distinct scopes correctly reuse stack slots, ensuring optimal stack usage.
This commit is contained in:
@@ -8,6 +8,8 @@ use std::rc::Rc;
|
||||
|
||||
struct StackAllocator {
|
||||
mapping: HashMap<u32, u32>,
|
||||
/// Physical slots that have been freed and can be reused.
|
||||
free_slots: Vec<u32>,
|
||||
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(|| {
|
||||
// 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<AnalyzedNode>]) -> Vec<VirtualId> {
|
||||
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<VirtualId>) {
|
||||
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);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user