Undefined behavior while popping custom Vec implementation

Hi, I'd been working on Implementing Vec chapter in Nomicon until I got stuck on a problem.

The problem is that when I try popping the first element, the function sometimes returns me a random value. As far as I've tested, the problem only occurs when there are more than one element in Vek.

Here are the Vek, pop() and the tests:

pub struct Vek<T> {
    ptr: NonNull<T>,
    len: usize,
    cap: usize,
}
pub fn pop(&mut self) -> Option<T> {
    if self.len == 0 {
        None
    } else {
        self.len -= 1;

        unsafe {
            Some(std::ptr::read(self.ptr.as_ptr().add(self.len)))
        }
    }
}
// As far as I've tried, this test always passes.
#[test]
fn pops_one_element() {
    let mut v: Vek<i32> = Vek::new();
    v.push(1);

    assert_eq!(v.pop(), Some(1));
    assert_eq!(v.pop(), None);
    assert_eq!(v.len, 0);
}

// This test sometimes fails, sometimes passes.
#[test]
fn pops_multiple_elements() {
    let mut v: Vek<i32> = Vek::new();
    v.push(1);
    v.push(2);
    v.push(3);

    assert_eq!(v.pop(), Some(3));
    assert_eq!(v.pop(), Some(2));
    assert_eq!(v.pop(), Some(1));
    assert_eq!(v.pop(), None);
    assert_eq!(v.len, 0);
}

Example error:

    Finished test [unoptimized + debuginfo] target(s) in 0.01s
     Running target/debug/deps/vek-eb5be6d87e20ccd1

running 6 tests
test tests::grows_as_doubles ... ok
test tests::initializes_empty_vek ... ok
test tests::grows_empty_to_one ... ok
test tests::pops_one_element ... ok
test tests::pushes_elements_and_grows_vec ... ok
test tests::pops_multiple_elements ... FAILED

failures:

---- tests::pops_multiple_elements stdout ----
thread 'tests::pops_multiple_elements' panicked at 'assertion failed: `(left == right)`
  left: `Some(-41287679)`,
 right: `Some(1)`', src/lib.rs:165:9
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace


failures:
    tests::pops_multiple_elements

test result: FAILED. 5 passed; 1 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

Full code:

use std::ptr::NonNull;
use std::alloc::{alloc, realloc, dealloc, Layout};
use std::ops::{Deref, DerefMut};

pub struct Vek<T> {
    ptr: NonNull<T>,
    len: usize,
    cap: usize,
}

impl<T> Vek<T> {
    pub fn new() -> Self {
        assert!(std::mem::size_of::<T>() > 0, "ZSTs aren't allowed.");
        Vek { ptr: NonNull::dangling(), len: 0, cap: 0 }
    }

    fn grow(&mut self) {
        if self.cap == 0 {
            unsafe {
                let layout = Layout::array::<T>(1).unwrap();
                let raw_ptr = alloc(layout) as *mut T;

                self.ptr = NonNull::new(raw_ptr).unwrap();
                self.cap = 1;
            }
        } else {
            let new_cap = self.cap * 2;
            let old_num_bytes = self.cap * std::mem::size_of::<T>();

            assert!(old_num_bytes <= isize::MAX as usize / 2, "Capacity overflow");

            unsafe {
                let raw_ptr = realloc(
                    self.ptr.as_ptr().cast(),
                    Layout::array::<T>(self.cap).unwrap(),
                    new_cap,
                ) as *mut T;

                self.ptr = NonNull::new(raw_ptr).unwrap();
                self.cap = new_cap;
           }
        }
    }

    pub fn push(&mut self, elem: T) {
        if self.len == self.cap { self.grow(); }
        
        unsafe {
            std::ptr::write(self.ptr.as_ptr().add(self.len), elem);
        }

        self.len += 1;
    }

    pub fn pop(&mut self) -> Option<T> {
        if self.len == 0 {
            None
        } else {
            self.len -= 1;

            unsafe {
                Some(std::ptr::read(self.ptr.as_ptr().add(self.len)))
            }
        }
    }
}

impl<T> Drop for Vek<T> {
    fn drop(&mut self) {
        if self.cap == 0 { return; }

        while let Some(_) = self.pop() {}

        unsafe {
            dealloc(
                self.ptr.as_ptr().cast(),
                Layout::array::<T>(self.cap).unwrap()
            );
        }
    }
}

impl<T> Deref for Vek<T> {
    type Target = [T];
    fn deref(&self) -> &[T] {
        unsafe {
            std::slice::from_raw_parts(self.ptr.as_ptr(), self.len)
        }
    }
}

impl<T> DerefMut for Vek<T> {
    fn deref_mut(&mut self) -> &mut [T] {
        unsafe {
            std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len)
        }
    }
}

P.S: I'm new to unsafe & this stuff and I want to learn more, so I'd be
grateful if you could explain things in detail :slight_smile:

Your realloc call needs the new size in bytes, like new_layout.size(), rather than new_cap.

3 Likes