dynvec/

raw.rs

use crate::region::Region;

use std::alloc::Layout;
use std::ops;

pub struct Header {
    /// A function that causes an item to be dropped.
    drop_fn: Option<fn(*mut u8)>,
}

/// A handle to an allocated object of a dynamic vector.
pub struct Handle<T> {
    ptr: *mut T,
}

impl<T> Clone for Handle<T> {
    fn clone(&self) -> Self {
        Self { ptr: self.ptr }
    }
}

impl<T> Copy for Handle<T> {}

impl<T> Handle<T> {
    /// Turns this `Handle` into a `RawHandle`.
    pub fn raw(self) -> RawHandle {
        RawHandle {
            ptr: self.ptr as *mut u8,
        }
    }
}

/// A handle to an allocated object of a dynamic vector that does
/// not own a `T`.
pub struct RawHandle {
    ptr: *mut u8,
}

impl Clone for RawHandle {
    fn clone(&self) -> Self {
        Self { ptr: self.ptr }
    }
}

impl Copy for RawHandle {}

impl<T> From<Handle<T>> for RawHandle {
    fn from(handle: Handle<T>) -> Self {
        handle.raw()
    }
}

impl RawHandle {
    /// Turns this `RawHandle` into a normal one, given it a type.
    ///
    /// # Safety
    /// The type `T` must be the one that created the raw handle.
    pub unsafe fn trust_type<T>(self) -> Handle<T> {
        Handle {
            ptr: self.ptr as *mut T,
        }
    }
}

/// A dynamic vector, able to store any kind of data.
#[derive(Default)]
pub struct RawDynVec<R: Region<Header>> {
    region: R,
}

impl<R: Region<Header>> Drop for RawDynVec<R> {
    fn drop(&mut self) {
        self.clear();
    }
}

impl<R: Region<Header>, T> ops::Index<Handle<T>> for RawDynVec<R> {
    type Output = T;

    fn index(&self, index: Handle<T>) -> &Self::Output {
        self.get(index)
    }
}

impl<R: Region<Header>, T> ops::IndexMut<Handle<T>> for RawDynVec<R> {
    fn index_mut(&mut self, index: Handle<T>) -> &mut Self::Output {
        self.get_mut(index)
    }
}

impl<R: Region<Header>> RawDynVec<R> {
    /// Creates a new `RawDynVec` with the given region.
    pub fn with_region(region: R) -> Self {
        Self { region }
    }

    /// Returns the number of items this vector stores.
    pub fn count(&self) -> usize {
        self.region.count()
    }

    /// Tries to insert a `T` into the vector and returns a Handle
    /// to this `T`. If the `T` cannot be allocated, it gets
    /// returned as the error.
    pub fn try_insert<T>(&mut self, item: T) -> Result<Handle<T>, T> {
        let layout = Layout::new::<T>();
        
        let drop_fn: Option<fn(*mut u8)> = if std::mem::needs_drop::<T>() {
            Some(|ptr: *mut u8| unsafe { std::ptr::drop_in_place(ptr as *mut T) })
        } else {
            None
        };
        
        match unsafe { self.insert_raw(layout, drop_fn) } {
            Ok(handle) => {
                unsafe {
                    handle.ptr.cast::<T>().write(item);
                    Ok(handle.trust_type::<T>())
                }
            },

            Err(()) => Err(item),
        }
    }

    /// Clears the whole vector, leaving it empty.
    pub fn clear(&mut self) {
        for (ptr, header) in self.region.deallocate_all() {
            if let Some(drop_fn) = header.drop_fn {
                // the allocated item needed to be dropped
                drop_fn(ptr)
            }
        }
    }

    /// Inserts a `T` into the vector and returns a Handle to this `T`.
    ///
    /// # Panics
    /// This function panics if the `T` could not be allocated.
    pub fn insert<T>(&mut self, item: T) -> Handle<T> {
        match self.try_insert(item) {
            Ok(handle) => handle,
            Err(_) => {
                panic!("Failed to allocate the item");
            }
        }
    }

    /// Allocates memory as described by the given layout. If the allocation
    /// is a success, a raw handle to the allocation is returned. In this case,
    /// it is up to the caller to initialize the allocated value because the
    /// vector now assumes it is initialized.
    /// 
    /// # Safety
    /// * The `drop_fn` function must cause the object to be dropped.
    /// * The allocated data must get initialized before it being dropped.
    pub unsafe fn insert_raw(
        &mut self,
        layout: Layout,
        drop_fn: Option<fn(*mut u8)>,
    ) -> Result<RawHandle, ()> {
        let header = Header {
            drop_fn,
        };

        match self.region.allocate(layout, header) {
            Ok(ptr) => {
                Ok(RawHandle { ptr, })
            }

            Err(_) => Err(()),
        }
    }

    /// Tries to remove the `T` located by the given handle. If the handle
    /// was invalid, `Err(())` is returned.
    pub fn remove<T>(&mut self, handle: Handle<T>) -> Result<T, ()> {
        if let Some(_) = self.region.deallocate(handle.ptr as *mut u8) {
            // the region successfuly deallocated the object
            // this means that we were able to initialize it before
            let item = unsafe { handle.ptr.read() };

            // we don't need to drop the item
            // this will be up to the caller
            Ok(item)
        } else {
            Err(())
        }
    }

    /// Tries to remove an item out of the vector. If the handle was
    /// invalid `Err(())` is returned.
    ///
    /// As the handle is raw, the item cannot be returned. But it is
    /// still dropped properly.
    pub fn remove_raw(&mut self, handle: RawHandle) -> Result<(), ()> {
        if let Some(header) = self.region.deallocate(handle.ptr) {
            if let Some(drop_fn) = header.drop_fn {
                // the item needed to be dropped
                drop_fn(handle.ptr);
            }

            Ok(())
        } else {
            Err(())
        }
    }

    /// Gets a pointer to the data located by the given handle.
    ///
    /// # Safety
    /// If the null pointer is returned, the handle was invalid.
    pub fn get_ptr<T>(&self, handle: Handle<T>) -> *const T {
        if self.region.has_allocated(handle.ptr as *mut u8) {
            handle.ptr
        } else {
            std::ptr::null()
        }
    }

    /// Gets a pointer to the data located by the given handle.
    ///
    /// # Safety
    /// If the null pointer is returned, the handle was invalid.
    pub fn get_mut_ptr<T>(&self, handle: Handle<T>) -> *mut T {
        if self.region.has_allocated(handle.ptr as *mut u8) {
            handle.ptr
        } else {
            std::ptr::null_mut()
        }
    }

    /// Gets a pointer to the data located by the given handle.
    ///
    /// # Safety
    /// If the null pointer is returned, the handle was invalid.
    pub fn get_ptr_raw(&mut self, handle: RawHandle) -> *const u8 {
        if self.region.has_allocated(handle.ptr) {
            handle.ptr
        } else {
            std::ptr::null()
        }
    }

    /// Gets a pointer to the data located by the given handle.
    ///
    /// # Safety
    /// If the null pointer is returned, the handle was invalid.
    pub fn get_mut_ptr_raw(&mut self, handle: RawHandle) -> *mut u8 {
        if self.region.has_allocated(handle.ptr) {
            handle.ptr
        } else {
            std::ptr::null_mut()
        }
    }

    /// Tries to get a reference to the data located by the given handle.
    pub fn try_get<T>(&self, handle: Handle<T>) -> Option<&T> {
        // if the handle is invalid, the null pointer is returned
        // this will transmute into `None`
        unsafe { std::mem::transmute(self.get_ptr(handle)) }
    }

    /// Tries to get a reference to the data located by the given handle.
    pub fn try_get_mut<T>(&mut self, handle: Handle<T>) -> Option<&mut T> {
        unsafe { std::mem::transmute(self.get_mut_ptr(handle)) }
    }

    /// Gets a reference to the data located by the given handle.
    ///
    /// # Panics
    /// If the handle is valid, this function panics.
    pub fn get<T>(&self, handle: Handle<T>) -> &T {
        self.try_get(handle).expect("Invalid handle")
    }

    /// Gets a reference to the data located by the given handle.
    ///
    /// # Panics
    /// If the handle is valid, this function panics.
    pub fn get_mut<T>(&mut self, handle: Handle<T>) -> &mut T {
        self.try_get_mut(handle).expect("Invalid handle")
    }
}

#[cfg(test)]
#[test]
fn properly_dropped() {
    use std::cell::Cell;
    use std::rc::Rc;

    use crate::region::block::Block;

    let drop_count = Rc::new(Cell::new(0u16));
    struct DropCheck(Rc<Cell<u16>>);
    impl Drop for DropCheck {
        fn drop(&mut self) {
            self.0.set(self.0.get() + 1);
        }
    }

    {
        let mut vec = RawDynVec::with_region(Block::new(2048));

        for _ in 0..100 {
            vec.insert(DropCheck(Rc::clone(&drop_count)));
        }

        // vec should be dropped now
    }

    assert_eq!(drop_count.get(), 100);
}

#[cfg(test)]
#[test]
fn normal_use() {
    use crate::region::block::Block;

    let mut vec = RawDynVec::with_region(Block::new(2048));

    let h_i32 = vec.insert(5i32);
    let h_vec = vec.insert(vec![1u8, 2, 3]);

    assert_eq!(vec[h_i32], 5i32);
    assert_eq!(vec[h_vec][1], 2u8);

    vec[h_vec].push(8);

    assert_eq!(vec.get(h_vec).last(), Some(&8));

    let other_vec = vec.remove(h_vec).unwrap();

    assert_eq!(&other_vec[..], &[1, 2, 3, 8][..]);

    assert_eq!(vec.try_get(h_vec), None);
}