ruststackmaybeuninit

Rust fast stack implementation without unnecessary memset?


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?


Solution

  • 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()
                }
            })
        }
    }