Rust fast stack implementation without unnecessary memset?

Question

I need a fast stack in Rust. Millions of these need to be created/destroyed per second and each of them need only a fixed depth. I'm trying to squeeze as much speed as I can. I came up with the following (basically textbook stack implementation):

const L: usize = 1024;
pub struct Stack {
    xs: [(u64, u64, u64, u64); L],
    sz: usize
}

impl Stack {
    pub fn new() -> Self {
        Self { xs: [(0, 0 ,0, 0); L], sz: 0 }
    }
    
    pub fn push(&mut self, item: (u64, u64, u64, u64)) -> bool {
        if (self.sz + 1) <= L {
            self.xs[self.sz] = item;
            self.sz += 1;
            true
        } else {
            false
        }
    }
    
    pub fn pop(&mut self) -> Option<(u64, u64, u64, u64)> {
        (self.sz > 0).then(|| {
            self.sz -= 1;
            self.xs[self.sz]
        })
    }
}

The problem is memset, which is unnecessary. So I tried to get rid of it:

    pub fn new2() -> Self {
        let xs = std::array::from_fn(|_| unsafe { MaybeUninit::uninit().assume_init() });
        Self { xs, sz: 0 }
    }

This gets rid of the memset, but now I have a warning:

  |
18 |         let xs = std::array::from_fn(|_| unsafe { MaybeUninit::uninit().assume_init() });
   |                                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   |                                                   |
   |                                                   this code causes undefined behavior when executed
   |                                                   help: use `MaybeUninit<T>` instead, and only call `assume_init` after initialization is done
   |
   = note: integers must not be uninitialized
   = note: `#[warn(invalid_value)]` on by default

If uninitialized integers cause undefined behavior, is it not possible to create this kind of stack, where the logic of the stack guarantees proper behavior and I avoid unnecessary memory operations?

Note that much of the hassle could be avoided altogether by using a plain `Vec`, or an [`arrayvec`](https://crates.io/crates/arrayvec). These ones are already prepared for fast insertion without initialization of its full capacity. https://stackoverflow.com/questions/38863315/how-to-perform-efficient-vector-initialization-in-rust https://stackoverflow.com/questions/31360993/what-is-the-proper-way-to-initialize-a-fixed-length-array — E_net4, Dec 20 '22 at 16:58
Are you sure that the initialization of the stack is a bottleneck? Because push and pop may very well be, but initialization is typically less hot. — Chayim Friedman, Dec 20 '22 at 17:02
@E_net4thecommentflagger I had a look at `arrayvec`. And interestingly, it does the [exact same thing as I tried](https://github.com/bluss/arrayvec/blob/0e131fd78101c9a4cd6c82f4005472d6fb92981a/src/arrayvec.rs#L80). Just that I didn't propogate the `MaybeUninit` all over. — , Dec 20 '22 at 17:04
@Ana Sure, this is how it is implemented, but there is a lot of benefit in using a ready, stable crate and not needing to worry about all the gory unsafe details. — Chayim Friedman, Dec 20 '22 at 17:05
@ChayimFriedman For something that's just a few lines of code, I'd rather avoid a dependency. I know this is a minority preference. — , Dec 20 '22 at 17:22
@Ana I typically agree, but when it comes down to avoiding unsafe code, I will add a dependency (and this is not a heavy one) even for one unsafe block. — Chayim Friedman, Dec 20 '22 at 17:41

score 1 · Accepted Answer · answered Dec 20 '22 at 16:55

You need to use MaybeUninit all the way through. Change your array to an array of MaybeUninits:

use std::mem::MaybeUninit;

const L: usize = 1024;
pub struct Stack {
    xs: [MaybeUninit<(u64, u64, u64, u64)>; L],
    sz: usize
}

// From standard library
// https://doc.rust-lang.org/stable/src/core/mem/maybe_uninit.rs.html#350-353
#[must_use]
#[inline(always)]
pub const fn uninit_array<const N: usize, T>() -> [MaybeUninit<T>; N] {
    // SAFETY: An uninitialized `[MaybeUninit<_>; LEN]` is valid.
    unsafe { MaybeUninit::<[MaybeUninit<T>; N]>::uninit().assume_init() }
}

impl Stack {
    pub fn new() -> Self {
        Self { xs: uninit_array(), sz: 0 }
    }
    
    pub fn push(&mut self, item: (u64, u64, u64, u64)) -> bool {
        if (self.sz + 1) <= L {
            self.xs[self.sz].write(item);
            self.sz += 1;
            true
        } else {
            false
        }
    }
    
    pub fn pop(&mut self) -> Option<(u64, u64, u64, u64)> {
        (self.sz > 0).then(|| {
            self.sz -= 1;
            // Safety: The value has been initialized
            unsafe {
                self.xs[self.sz].assume_init()
            }
        })
    }
}

Ohh ... I misunderstood the meaning of MaybeUninit. Thanks. – Dec 20 '22 at 16:59 — , Dec 20 '22 at 16:59

Rust fast stack implementation without unnecessary memset?

1 Answers1