Demilade Sonuga's blog

All posts

Vec II

2023-04-11

In this post, we're going to implement a vector.

Implementation

The role of the Vec is to allocate space for a variable-length sequence of data items on the heap, store that data, allocate more heap space when the sequence becomes too large for the initially allocated space, and deallocate that data when the Vec is no longer needed.

To do this, we need to keep track of just four pieces of data: the length (number of data items currently stored in the allocated heap space), the capacity (number of data items that can be stored in the allocated heap space), a pointer to the allocated heap space, and a pointer to the allocator.

Create a new file vec.rs and type this in:

use crate::allocator::Allocator;

// A growable array on the heap
pub struct Vec<T: Clone> {
    len: usize,
    capacity: usize,
    start_ptr: *mut T,
    allocator: *mut dyn Allocator
}

If you read the Vec chapters in the rustonomicon or if you looked at the Vec implementation in std or alloc, you'll notice that it's different from this. You might even say that this is a naive way of doing this (it is), but we need to remember that we're not building a general Vec for anyone to use. It has a very special purpose in our game, so we don't need to give it any special abilities. The simplest implementation that can do what we need it to do is good enough.

Again, we're using generics here to tell the compiler that the Vec can hold data items of any type. The Clone in Vec<T: Clone> restricts the types that vectors can hold. It tells the compiler that the vector can hold, not just any type, but rather, any type that implements Clone. This is necessary because if an item in a vector doesn't implement Clone, how will we be able to clone the vector?

Functions we're going to use to create and handle the Vec:

  1. with_capacity for creating a vector with an initial amount of space on the heap.
  2. push for appending a data item to the end of a vector.
  3. pop for removing and returning the last data item in a vector.
  4. try_pop: does the same thing as pop, but returns an Option, None in the case of an empty vector.
  5. remove for removing an item in the vector at a particular index.
  6. len for returning the current length of the vector.
  7. capacity for returning the current capacity of the vector.
  8. as_ptr for returning the start pointer of the vector's heap data.
  9. iter and iter_mut for creating iterators over the vector's contents.

We're also going to implement a bunch of traits: Drop, Index, IndexMut, PartialEq, and Clone.

If you haven't done so before, implement all the functions before moving on.

Starting from with_capacity. To create a new Vec, we'll need the initial capacity and a pointer to an allocator to allocate the required capacity.

impl<T: Clone> Vec<T> {

    // Creates a vector with the stated capacity
    pub fn with_capacity(capacity: usize, allocator: *mut dyn Allocator) -> Vec<T> {
        
    }
}

Now, in the function, we'll first attempt to allocate enough space to hold capacity items of type T. If it's successful, a new Vec is created and returned. If it's not successful, a panic with an error message is the result.

use core::mem;

impl<T: Clone> Vec<T> {

    pub fn with_capacity(capacity: usize, allocator: *mut dyn Allocator) -> Vec<T> {
        // Allocate heap memory to hold `capacity` items of type `T`
        let alloc_result = unsafe { (*allocator).alloc(mem::size_of::<T>() * capacity, mem::align_of::<T>()) };
        match alloc_result {
            // Heap allocation succeeded
            Some(ptr) => Vec {
                len: 0,
                capacity,
                start_ptr: ptr as *mut T,
                allocator
            },
            // Heap allocation failed
            None => panic!("No enough space on the heap")
        }
    }
}

The next item on our list is the push. When pushing there are two cases to consider. The first is when there is enough space in the vector to accommodate the new data item, that is, capacity > len, so adding another item will not exceed the vector's capacity. But in the second case, when there isn't enough space in the vector to hold the data item, the vector needs to reallocate memory to accommodate the new item. Like anything else in the programming world, there is definitely more than one way of dealing with this situation, but the approach we're going to use here is very common:

  1. Allocate a new memory chunk double the size of the current capacity.
  2. Move the vector's data from the former location to the newly allocated location.
  3. Deallocate the previous location.

Now, the vector has enough space to hold another item.

For the first case:

impl<T: Clone> Vec<T> {
    // Add an item to the end of the vector
    pub fn push(&mut self, item: T) {
        // There is enough space for the new item
        if self.capacity > self.len {
            unsafe { 
                let ptr_to_next_position = self.start_ptr.offset(self.len as isize);
                ptr_to_next_position.write(item);
            }
            self.len += 1;
        }
    }
}

The code is straightforward. Put the item at the end. Increase the length. But there are some subtleties involved. For one, notice that the code says ptr_to_next_position.write(item) and not *ptr_to_next_position = item. We have to be very careful here because if we aren't, we'll reap the undefined' behaviors that come with messing around with unsafe.

To understand why the write function was used here, you need to take a look at the documentation page on Primitive Type pointer. According to these docs, *ptr = data results in a call to drop on the old value. What this means is that ptr will first be dereferenced. Whatever it dereferences to will be interpreted as what its type says it's pointing to. Then the drop implementation will be called.

From this information, you should now see why we can't just say *ptr_to_next_position = item. This is because ptr_to_next_position doesn't yet point to a valid item. We've never stored anything there. But if we do this, the pointer will be dereferenced as if we've stored something there before and a drop implementation will be attempted, but from where can you get a drop implementation if there is no value there in the first place?

The result will be undefined behavior, headache, and frustration all because of one innocent-looking line of code. This is why we use the write function, because write doesn't call drop first before writing the data to the location the pointer is pointing to.

To make this more clear, you can say that *ptr_to_next_position = item is just a fancy way of saying:

(*ptr_to_next_position).drop()
ptr_to_next_position.write(item);

In the second case, we need to allocate double the capacity, move data into the new location, deallocate the previous space, and update the vector's fields to reflect the new heap memory and capacity.

impl<T: Clone> Vec<T> {
    pub fn push(&mut self, item: T) {
        if self.capacity > self.len {
            /* ...Others */
        // NEW:
        } else {
            // Allocate double the capacity
            let new_size = self.capacity * 2;
            let old_size = self.capacity;
            let old_start_ptr = self.start_ptr as *mut u8;
            let new_start_ptr = unsafe {
                (*self.allocator).alloc(mem::size_of::<T>()  * new_size, mem::align_of::<T>())
            }.expect("No enough space on the heap.") as *mut T;
            // Move data into the new location
            for i in 0..self.len {
                unsafe {
                    let val = self.start_ptr.offset(i as isize).read();
                    new_start_ptr.offset(i as isize).write(val);
                }
            }
            // Deallocate the previous heap memory
            unsafe { (*self.allocator).dealloc(old_start_ptr, old_size * mem::size_of::<T>()) };
            // Update the fields to reflect the new capacity and heap memory
            self.capacity = new_size;
            self.start_ptr = new_start_ptr as *mut T;
        }
    }
}

The next function on our list is pop. pop is very straightforward. If there the length of the vector is 0 (no items to pop), a panic occurs. Else, the len field is decreased by 1 and the last item is returned.

impl<T: Clone> Vec<T> {
    // Removes and returns the last item in the vector
    // Panics when there are no items in the vector
    pub fn pop(&mut self) -> T {
        if self.len == 0 {
            panic!("No items to pop");
        }
        self.len -= 1;
        unsafe { self.start_ptr.offset(self.len as isize).read() }
    }
}

try_pop is just like pop, but instead of panicking on an empty vector, a None is returned.

impl<T: Clone> Vec<T> {
    // Removes and returns the last item in the vector
    // When there are no items in the vector, a `None` is returned
    pub fn try_pop(&mut self) -> Option<T> {
        if self.len == 0 {
            None
        } else {
            Some(self.pop())
        }
    }
}

remove is a bit tricky. It takes an index and removes the item at that index from the vector.

impl<T: Clone> Vec<T> {
    // Removes the item at index `idx` and returns it
    // Panics when `idx` is an invalid index
    pub fn remove(&mut self, idx: usize) -> T {

    }
}

We first need to verify that the index is valid. After doing that, we need to retrieve the value (carefully, so we don't mistakenly drop it), then shift the remaining elements in the vector to fill up the space left by the removed element. For example, consider this idealized vector:

index:  0    1       2              3       4
vector: "I", "love", "programming", "with", "Rust"

If remove is called for this vector with an argument of 2, we first verify that 2 is a valid index. Since 2 is valid, we retrieve its value (without dropping it):

retrieved_value: "programming"
index:  0    1       2              3       4
vector: "I", "love",              , "with", "Rust"

Then we have to shift the elements at indexes 3 and 4 into the positions at 2 and 3, respectively. Why? So that the vector will keep being a vector.

retrieved_value: "programming"
index:  0    1       2       3
vector: "I", "love", "with", "Rust"

In Rust:

impl<T: Clone> Vec<T> {
    pub fn remove(&mut self, idx: usize) -> T {
        // First, check if the index is valid
        if idx >= self.len {
            panic!("Invalid index");
        }
        // Retrieve the value at the index `idx`
        let value = unsafe { self.start_ptr.offset(idx as isize).read() };
        // Shift the elements that come after the removed value
        // to cover the hole (without dropping them)
        for i in idx + 1..self.len {
            let i = i as isize;
            unsafe {
                let val = self.start_ptr.offset(i).read();
                self.start_ptr.offset(i - 1).write(val);
            }
        }
        // Update the `len` field to reflect the new vector state
        self.len -= 1;
        value
    }
}

Functions len, capacity, and as_ptr are pretty easy since they're just returning reading and returning values.

impl<T: Clone> Vec<T> {
    // Returns the number of items in the vector
    pub fn len(&self) -> usize {
        self.len
    }

    // Returns the capacity of the vector
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    // Returns the pointer to the vector data
    pub fn as_ptr(&self) -> *const T {
        self.start_ptr
    }
}

As for the last functions: iter and iter_mut, we need to create iterators over the vector data. iter for iterating over immutable references of the data and iter_mut for iterating over mutable ones.

We can create our own struct and implement Iterator for it (do that if you want to) or we can use the utilities offered to use by core::slice. A slice has iter and iter_mut functions that do exactly what we want and core::slice has functions that we can use to create a slice from our vector data.

impl<T: Clone> Vec<T> {
    // Creates a new iterator over the references of the vector's elements
    pub fn iter(&self) -> core::slice::Iter<T> {
        unsafe { core::slice::from_raw_parts(self.start_ptr as *const T, self.len) }
            .iter()
    }

    // Creates a new iterator over mutable references of the vector's elements
    pub fn iter_mut(&mut self) -> core::slice::IterMut<T> {
        unsafe { core::slice::from_raw_parts_mut(self.start_ptr, self.len) }
            .iter_mut()
    }
}

iter first calls core::slice's from_raw_parts function to create an immutable slice out of the vector's elements. The slice's iter function then does the rest. The same thing is done with iter_mut. The from_raw_parts_mut function creates a mutable slice of the vector's elements. The slice's iter_mut function does the rest.

What comes next is our Drop implementation. We first have to drop all the values in the vector, then deallocate the memory.

impl<T: Clone> Drop for Vec<T> {
    fn drop(&mut self) {
        use core::ptr;
        unsafe {
            ptr::drop_in_place(ptr::slice_from_raw_parts_mut(self.start_ptr, self.len));
            (*self.allocator).dealloc(self.start_ptr as *mut u8, self.capacity * mem::size_of::<T>());
        };
    }
}

We first create a mutable slice from the vector's elements, then we pass that slice as an argument to ptr's drop_in_place. drop_in_place then does the job of dropping all the elements.

The next traits to implement are Index and IndexMut. These traits are the reasons a vector can be indexed like vector[0].

use core::mem;
use crate::allocator::Allocator;
use core::ops::{Index, IndexMut}; // NEW

impl<T: Clone> Index<usize> for Vec<T> {
    type Output = T;

    fn index(&self, idx: usize) -> &Self::Output {
        assert!(idx < self.len);
        unsafe { &*self.start_ptr.offset(idx as isize) }
    }
}

impl<T: Clone> IndexMut<usize> for Vec<T> {
    fn index_mut(&mut self, idx: usize) -> &mut Self::Output {
        assert!(idx < self.len);
        unsafe { &mut *self.start_ptr.offset(idx as isize) }
    }
}

If you aren't familiar with the Index and IndexMut traits, check the references. As you can see, Index is for returning immutable references to the item we're retrieving, while IndexMut is for mutable references.

The start pointer's offset function is invoked to get a pointer to the exact element we're interested in (the one at index idx). That pointer is then dereferenced and a safe reference is obtained for the underlying data.

With these implementations, our Vec can now be indexed with a usize.

As for PartialEq, we can't just implement it for any Vec. We can only implement PartialEq for Vecs containing elements that can be compared. This is simply resolved by adding a new trait bound.

impl<T: PartialEq + Clone> PartialEq<Vec<T>> for Vec<T> {
    fn eq(&self, other: &Vec<T>) -> bool {
        if self.len() != other.len() {
            return false;
        }
        for (val1, val2) in self.iter().zip(other.iter()) {
            if val1 != val2 {
                return false;
            }
        }
        true
    }
}

As you can see, the PartialEq in T: PartialEq + Clone means that the compiler should implement this trait only for vectors of items that can be compared (that already implement PartialEq).

We use the Iterator's zip function to go through the two vector elements at once. If one vector is shorter than the other, the iteration will end at the point that the shorter vector ends. That's the reason why we need that self.len() == other.len() condition.

If you aren't familiar with the Iterator trait and the stuff it has to offer, you should really check it out and experiment with them. The link to the docs is in the references.

The last trait on our list is Clone. This should result in the creation of an identical vector. This is also the reason why we made Clone a trait bound in Vec's definition.

impl<T: Clone> Clone for Vec<T> {
    fn clone(&self) -> Self {
        let mut new_vec = Vec::with_capacity(self.capacity, self.allocator);
        self
            .iter()
            .for_each(|val| new_vec.push(val.clone()));
        new_vec
    }
}

The effect of this implementation is to create a new vector with the same capacity and then clone all the elements of the original into the new.

That's it for our implementation. In the next post, we'll be testing Vec.

Take Away

For the full code, go to the repo

In The Next Post

We'll be testing Vec

References

  • Primitive Type pointer documentation: https://doc.rust-lang.org/core/primitive.pointer.html
  • Index documentation: https://doc.rust-lang.org/core/ops/trait.Index.html
  • IndexMut documentation: https://doc.rust-lang.org/core/ops/trait.IndexMut.html
  • Documentation related to iterators: https://doc.rust-lang.org/core/iter/index.html

← Previous Post

Vec I

Next Post →

Vec III