Interpreter spill stack on the process stack & trait::Allocator for stack allocations

Oct 10, 2026
2503 Words
12 Minute read

Tags:

I’m currently still working on my scripting language, purple-garden, and wanted to port all stages of the compilation pipeline over to the new allocator trait (for a list of reasons, for instance i use scratch storage in all stages for temporary allocs and would like to both reuse and optimise this memory area for locality and cheap allocs) one of these zig style inspired allocators is the StackAllocator which ill explain in this article and use to optimise the register spilling and restore to the virtual machine stack (the others are PageAlloc, BumpAlloc<A> and MetricAllocator<A>).

These changes dont really affect the general purpose performance of purple garden anymore, since most of the speed is squeezed out of the runtime via the JIT (only x86 for now)

Some prerequisites

trait::Allocator

The trait::Allocator proposal was recently stabilised in rust version 1.100 which introduces a new unsafe trait and a few functions taking these allocators, for instance {Vec,Box,HashMap,HashSet}::new_in(). These functions then allocate their memory into the passed in allocator and thus can now live in whatever memory region the allocator is backed by:

RUST

 1pub unsafe trait Allocator {
 2    fn allocate(
 3        &self,
 4        layout: Layout,
 5    ) -> Result<NonNull<[u8]>, AllocError>;
 6
 7    unsafe fn deallocate(
 8        &self,
 9        ptr: NonNull<u8>,
10        layout: Layout,
11    );
12
13    fn allocate_zeroed(
14        &self,
15        layout: Layout,
16    ) -> Result<NonNull<[u8]>, AllocError>
17
18    unsafe fn grow(
19        &self,
20        ptr: NonNull<u8>,
21        old_layout: Layout,
22        new_layout: Layout,
23    ) -> Result<NonNull<[u8]>, AllocError>
24
25    unsafe fn grow_zeroed(
26        &self,
27        ptr: NonNull<u8>,
28        old_layout: Layout,
29        new_layout: Layout,
30    ) -> Result<NonNull<[u8]>, AllocError>
31
32    unsafe fn shrink(
33        &self,
34        ptr: NonNull<u8>,
35        old_layout: Layout,
36        new_layout: Layout,
37    ) -> Result<NonNull<[u8]>, AllocError>
38}

Purple garden virtual machine

Purple gardens interpreter and thus the virtual machine are register based (for performance reasons) and therefore may run out of space in the register. This is indicated by the bytecode register allocator for both running out of register and also extending values lifetimes over functions, for instance consider the humble fib function:

GARDEN

#! Compute Fibonacci numbers with the classic doubly-recursive definition.
fn fib(n:Int) Int {
    match {
        n == 0 { 0 }
        n == 1 { 1 }
        { fib(n - 1) + fib(n - 2) }
    }
}

And its garden disassembly:

ASM

 1; purple-garden --no-jit -D examples/fib.garden
 200000000 <fib>:
 3  ; 8: fn fib(n:Int) Int {
 4  0000:    push r1
 5  ; 10: n == 0 { 0 }
 6  0001:    jmpne_imm r0, #0, 0006 <fib.bb_0006>
 7  0002:    load_imm r1, #0
 8  0003:    mov r0, r1
 9  0004:    pop r1
10  0005:    ret
11
1200000006 <fib.bb_0006>:
13  ; 11: n == 1 { 1 }
14  0006:    jmpne_imm r0, #1, 000b <fib.bb_000b>
15  0007:    load_imm r1, #1
16  0008:    mov r0, r1
17  0009:    pop r1
18  000a:    ret
19
200000000b <fib.bb_000b>:
21  ; 12: { fib(n - 1) + fib(n - 2) }
22  000b:    isub_imm r1, r0, #1
23  000c:    push r0
24  000d:    mov r0, r1
25  000e:    call 0000 <fib>
26  000f:    mov r1, r0
27  0010:    pop r0
28  0011:    isub_imm r0, r0, #2
29  0012:    call 0000 <fib>
30  0013:    iadd r0, r1, r0
31  0014:    pop r1
32  0015:    ret

In this example the function prologue (push r1) and epilogue (pop r1) serve the purpose of saving and restoring callee saved registers the recursive calls would clobber.

Info - Purple garden calling convention excursion

Registers are split into two zones (similar to x86 and aarch64):

r0..r{n} is considered the argument zone, while r0 is also the return slots, these are caller saved so a caller with a live value in this range must spill and restore around the call
r{n+1}..r{max} are the callee saved registers, meaning the callee doesnt clobber them

The fib example would be better by putting constants into better fitting registers, for instance r0 instead of the indirection through r1 and only pushing in the branches really clobbering r1, but this is zukunftsmusik, as the germans say :)

Previous Vm stack, {Push,Pop}{2,3} and their handlers

The stack interaction opcodes for the vm are defined as:

RUST

1Push { src: u8, },
2Push2 { a: u8, b: u8, },
3Push3 { a: u8, b: u8, c: u8, },
4Pop { dst: u8, },
5Pop2 { a: u8, b: u8, },
6Pop3 { a: u8, b: u8, c: u8, },

The simplest stack i could think of when implementing the virtual machine was using Vec<Value> for Vm::spilled, preallocating it with 4096 values and pushing and popping with unchecked mut writes and reads:

RUST

 1Op::Push { src } => {
 2    // whole handler basically
 3    // equivalent to self.spilled.push(*r!(src))
 4
 5    let len = self.spilled.len();
 6
 7    if std::hint::unlikely(
 8        len + 1 > self.spilled.capacity()) {
 9        self.spilled.reserve(1);
10    }
11
12    unsafe {
13        // use mut pointer, advance to length
14        let dst = self.spilled.as_mut_ptr().add(len);
15        // write value in register src into spilled[len]
16        dst.write(*r!(src));
17        self.spilled.set_len(len + 1);
18    }
19}
20
21// ...
22
23Op::Pop { dst } => unsafe {
24    // whole handler basically equivalent to
25    // *r_mut!(dst) = self.spilled.pop();
26
27    let len = self.spilled.len();
28    let ptr = self.spilled.as_ptr();
29    debug_assert!(len >= 1);
30    r_mut!(dst) = ptr.add(len - 1).read();
31    self.spilled.set_len(len - 1);
32}

Then i also added the fused {Push,Pop}{2,3} merging two and three stack interactions into a single operation without bounds check, for instance Pop3:

RUST

1Op::Pop3 { a, b, c } => unsafe {
2    let len = self.spilled.len();
3    let ptr = self.spilled.as_ptr();
4    debug_assert!(len >= 3);
5    r_mut!(a) = ptr.add(len - 1).read();
6    r_mut!(b) = ptr.add(len - 2).read();
7    r_mut!(c) = ptr.add(len - 3).read();
8    self.spilled.set_len(len - 3);
9}

StackAllocator

Technically just an allocator allocating into an already allocated buffer that just happens to be on the stack :O (I guess naming this BufferAllocator or something similar would be more useful, but im using it for allocating on the stack so its fine). For instance with a buffer allocated on the stack by (us) the caller, like so:

RUST

1let mut buf = [MaybeUninit::<u8>::uninit(), 1024];
2let alloc = StackAlloc::new(&mut spillstack);
3let mut vec_on_the_stack = Vec::new_in(&alloc);
4vec_on_the_stack.push(787);
5vec_on_the_stack.push(161);
6vec_on_the_stack.push(256);
7let box_on_the_top = Box::new_in(0xDEADAFFE, &alloc)

StackAlloc doesnt consume or use the handed in buffer other than keeping it around as a marker for lifetime purposes so the compiler makes sure the data doesnt outlive our data living in it. It then has a base, named block, a cap for bounds checks and top for advancing.

RUST

 1pub struct StackAlloc<'buf> {
 2    /// base ptr
 3    block: NonNull<u8>,
 4    /// buffer size
 5    cap: usize,
 6    /// next free byte
 7    top: Cell<usize>,
 8    /// marker for lifetime >:(
 9    _buf: PhantomData<&'buf mut [MaybeUninit<u8>]>,
10}

The next chapter will explain the Allocator trait impl, but at a high level the allocator works by:

find the start, which is block+top; rounded up to alignment (16 bytes, see “Tip” below)
find the end of the allocation; requested size rounded up to alignment, added to start
check if enough space in buffer; end > cap
advance top to end
return block+start, sliced to request size to omit alignment

Tip - Why alignment is 16 bytes

StackAllocator uses a 16 byte alignment to fix the issue of multiple alignments leaving padding behind, for instance: 3 bytes at offset 0, then 8 bytes aligned to 8 at offset 8. Freeing the 8 bytes moves top back to 8, but the 3 bytes end at 3, so they never count as the top allocation again and are lost until the allocator is reset. With rounding up to 16 byte alignment, every allocation aligned up to 16 starts right at top and frees in reverse work.

The methods supporting the virtual machine stack dont use this abstraction, since all values in the vm are a uniform 8 bytes.

Impl trait::Allocator

RUST

 1impl<'buf> StackAlloc<'buf> {
 2    pub fn new(buf: &'buf mut [MaybeUninit<u8>]) -> Self {
 3        let skip = buf.as_ptr().align_offset(UPROUND).min(buf.len());
 4        let block = NonNull::from(&mut buf[skip..]).cast::<u8>();
 5        Self {
 6            block,
 7            cap: (buf.len() - skip) / UPROUND * UPROUND,
 8            top: Cell::new(0),
 9            _buf: PhantomData,
10        }
11    }
12
13    #[must_use]
14    pub fn depth(&self) -> usize {
15        self.top.get()
16    }
17
18    pub fn reset(&mut self) {
19        self.top.set(0);
20    }
21
22    fn offset(&self, ptr: NonNull<u8>) -> usize {
23        ptr.as_ptr() as usize - self.block.as_ptr() as usize
24    }
25
26    fn is_top(&self, ptr: NonNull<u8>, layout: Layout) -> bool {
27        self.offset(ptr) + granules(layout.size()) == self.top.get()
28    }
29}
30
31// singular buffer borrow, so sending is safe
32unsafe impl Send for StackAlloc<'_> {}

depth, reset, offset and is_top are used in the allocator trait impl and the vm stack methods:

RUST

 1
 2pub const UPROUND: usize = 16;
 3
 4#[inline(always)]
 5fn granules(size: usize) -> usize {
 6    size.next_multiple_of(UPROUND)
 7}
 8
 9unsafe impl Allocator for StackAlloc<'_> {
10    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
11        let base = self.block.as_ptr() as usize;
12        let start = (base + self.top.get()).next_multiple_of(layout.align()) - base;
13        let end = start
14            .checked_add(granules(layout.size()))
15            .filter(|&end| end <= self.cap)
16            .ok_or(AllocError)?;
17        self.top.set(end);
18        let ptr = unsafe { self.block.add(start) };
19        Ok(NonNull::slice_from_raw_parts(ptr, layout.size()))
20    }
21
22    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
23        if self.is_top(ptr, layout) {
24            self.top.set(self.offset(ptr));
25        }
26    }
27
28    unsafe fn grow(
29        &self,
30        ptr: NonNull<u8>,
31        old: Layout,
32        new: Layout,
33    ) -> Result<NonNull<[u8]>, AllocError> {
34        if self.is_top(ptr, old) && (ptr.as_ptr() as usize).is_multiple_of(new.align()) {
35            let end = self
36                .offset(ptr)
37                .checked_add(granules(new.size()))
38                .filter(|&end| end <= self.cap)
39                .ok_or(AllocError)?;
40            self.top.set(end);
41            return Ok(NonNull::slice_from_raw_parts(ptr, new.size()));
42        }
43        let moved = self.allocate(new)?;
44        unsafe {
45            ptr.copy_to_nonoverlapping(moved.cast(), old.size());
46            self.deallocate(ptr, old);
47        }
48        Ok(moved)
49    }
50
51    unsafe fn grow_zeroed(
52        &self,
53        ptr: NonNull<u8>,
54        old: Layout,
55        new: Layout,
56    ) -> Result<NonNull<[u8]>, AllocError> {
57        let block = unsafe { self.grow(ptr, old, new)? };
58        unsafe {
59            block
60                .cast::<u8>()
61                .add(old.size())
62                .write_bytes(0, new.size() - old.size());
63        }
64        Ok(block)
65    }
66
67    unsafe fn shrink(
68        &self,
69        ptr: NonNull<u8>,
70        old: Layout,
71        new: Layout,
72    ) -> Result<NonNull<[u8]>, AllocError> {
73        if !(ptr.as_ptr() as usize).is_multiple_of(new.align()) {
74            let moved = self.allocate(new)?;
75            unsafe {
76                ptr.copy_to_nonoverlapping(moved.cast(), new.size());
77                self.deallocate(ptr, old);
78            }
79            return Ok(moved);
80        }
81        if self.is_top(ptr, old) {
82            self.top.set(self.offset(ptr) + granules(new.size()));
83        }
84        Ok(NonNull::slice_from_raw_parts(ptr, new.size()))
85    }
86}

StackAllocator::{push,pop,push_all,pop_all}

Now, for the virtual machine spill stack interaction, pushing and popping is necessary, for this i decided to implement these on the allocator itself, noticeably these dont align up to 16, they instead align to T:

RUST

 1#[inline(always)]
 2pub fn push<T: Copy>(&self, value: T) -> bool {
 3    let top = self.top.get();
 4    debug_assert!(top.is_multiple_of(align_of::<T>()));
 5    let end = top + size_of::<T>();
 6    if std::hint::unlikely(end > self.cap) {
 7        return false;
 8    }
 9    unsafe { self.block.add(top).cast::<T>().write(value) };
10    self.top.set(end);
11    true
12}
13
14#[inline(always)]
15pub unsafe fn pop<T: Copy>(&self) -> T {
16    let top = self.top.get() - size_of::<T>();
17    self.top.set(top);
18    unsafe { self.block.add(top).cast::<T>().read() }
19}
20
21#[inline(always)]
22pub fn push_all<T: Copy, const N: usize>(&self, values: [T; N]) -> bool {
23    let top = self.top.get();
24    debug_assert!(top.is_multiple_of(align_of::<T>()));
25    let end = top + N * size_of::<T>();
26    if std::hint::unlikely(end > self.cap) {
27        return false;
28    }
29    let dst = unsafe { self.block.add(top).cast::<T>() };
30    for (i, value) in values.into_iter().enumerate() {
31        unsafe { dst.add(i).write(value) };
32    }
33    self.top.set(end);
34    true
35}
36
37#[inline(always)]
38pub unsafe fn pop_all<T: Copy, const N: usize>(&self) -> [T; N] {
39    let top = self.top.get() - N * size_of::<T>();
40    self.top.set(top);
41    let src = unsafe { self.block.add(top).cast::<T>() };
42    std::array::from_fn(|i| unsafe { src.add(i).read() })
43}

Now: Vm stack with actual stack space

So now given these helper methods I can replace the previous handlers with an actual stack based spill stack, first allocating a stack buffer and creating the allocator in Vm::run:

RUST

1// 262144 B => 256KiB to spill results in /8 => 32768 vm values
2pub const SPILL_STACK: usize = 256 << 10;
3
4let mut spill_stack = [MaybeUninit::<u8>::uninit(); SPILL_STACK];
5let spilled = StackAlloc::new(&mut spill_stack);

This is then used in all stack op code handlers:

RUST

 1Op::Push { src } => unsafe {
 2    if std::hint::unlikely(!spilled.push(*r!(src))) {
 3        return Err(Anomaly::StackOverflow { pc });
 4    }
 5},
 6Op::Push2 { a, b } => unsafe {
 7    if std::hint::unlikely(!spilled.push_all([*r!(a), *r!(b)])) {
 8        return Err(Anomaly::StackOverflow { pc });
 9    }
10},
11Op::Push3 { a, b, c } => unsafe {
12    if std::hint::unlikely(!spilled.push_all([*r!(a), *r!(b), *r!(c)])) {
13        return Err(Anomaly::StackOverflow { pc });
14    }
15},
16Op::Pop { dst } => unsafe {
17    r_mut!(dst) = spilled.pop::<Value>();
18},
19Op::Pop2 { a, b } => unsafe {
20    let [below, top] = spilled.pop_all::<Value, 2>();
21    r_mut!(a) = top;
22    r_mut!(b) = below;
23},
24Op::Pop3 { a, b, c } => unsafe {
25    let [bottom, below, top] = spilled.pop_all::<Value, 3>();
26    r_mut!(a) = top;
27    r_mut!(b) = below;
28    r_mut!(c) = bottom;
29},

BENCHMARKS

For the commit 1296116 this article is about, see bench/cross for the benchmark source code.

Spill stack changes

Benched on over 10 runs after 3 warmup runs, interpreter only (–no-jit), release builds with -C target-cpu=native on an AMD Ryzen 7 3700X. Before is the parent of commit 1296116, after is 1296116:

bench	before (ms)	after (ms)	speedup
ackermann	162.9	129.1	1.26x / 20.8%
collatz	104.3	89.0	1.17x / 14.7%
fib	62.5	55.4	1.13x / 11.4%
gcd_sum	93.9	82.7	1.14x / 11.9%
keyword_table	19.3	18.0	1.07x / 6.5%
mandelbrot	73.4	60.7	1.21x / 17.4%
n_body	28.7	25.1	1.14x / 12.6%
primes_count	50.9	46.4	1.10x / 8.8%
tak	322.5	267.5	1.21x / 17.0%

So we are 6-20% faster for these benchmarks, thats a good results for such a non invasive change :).

VS with JIT enabled and against bun, luajit:

This is meant as a quick compare to how fast garden is when the jit is enabled vs when the jit is disabled (see previous table). I removed python from this table since it is so immensly slower than bun, luajit and garden:

The jit doesnt yet support doubles, so it will be even quicker in the future O.O

I have cross language benchmarks available at xnacly.github.io/purple-garden/, the table below is from commit da1ad03:

workload	garden (ms)	bun (ms)	luajit (ms)	garden vs fastest
ackermann	39.3	40.8	39.8	fastest
collatz	20.7	22.2	14.1	1.47x slower vs luajit
fib	17.9	16.1	9.61	1.86x slower vs luajit
gcd_sum	27.4	33.7	45.1	fastest
keyword_table	16.6	762	28.4	fastest
mandelbrot	89.1	14.2	6.45	13.81x slower vs luajit
n_body	49.7	16.4	5.02	9.90x slower vs luajit
primes_count	7.96	17.7	6.09	1.31x slower vs luajit
tak	58.2	53.0	55.9	1.10x slower vs bun

The website showing the results is embedded below (meaning this will update for each new commit, they are measured with jit enabled though):