allocal/

lib.rs

use std::ptr::NonNull;
use std::any::TypeId;
use std::marker::PhantomData;

/// A structure that allows the allocation of multiple objects of different types
/// in a contiguous region of the memory.
pub struct Allocal {
    /// A pointer to the allocated data.
    data: NonNull<u8>,

    /// The number of bytes that are currently allocated.
    cap: usize,

    /// A vector that keeps tarck of all the allocated objects.
    /// This vector is sorted by offset.
    locations: Vec<LocationInner>,
}

/// Locates an allocated object within a `Allocal`.
#[derive(Clone)]
pub struct LocationInner {
    /// The offset of the object, in bytes, from the begining of the allocated data.
    offset: usize,
    /// The size, in bytes, of the object.
    size: usize,
    /// A function that causes the object to be dropped.
    drop_fn: Option<fn(*mut u8)>,
}

pub struct Location<T> {
    offset: usize,
    _m: PhantomData<*const T>,
}

impl<T> Clone for Location<T> {
    fn clone(&self) -> Self {
        Self {
            offset: self.offset,
            _m: PhantomData,
        }
    }
}

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

impl<T: 'static, U: 'static> PartialEq<Location<U>> for Location<T> {
    #[inline(always)]
    fn eq(&self, other: &Location<U>) -> bool {
        TypeId::of::<T>() == TypeId::of::<U>() && self.offset == other.offset
    }
}

impl Drop for Allocal {
    fn drop(&mut self) {
        use std::alloc::{dealloc, Layout};

        // 1. Drop the allocated objects
        self.clear();

        // 2. Deallocate the data
        let layout = unsafe { Layout::from_size_align_unchecked(self.cap, 1) };
        unsafe { dealloc(self.data.as_ptr(), layout) };
    }
}

impl Allocal {
    /// Allocates a new `Allocal` with the given capacity.
    /// 
    /// # Panics
    /// This function panics if the data is unable to be allocated.
    pub fn with_capacity(capacity: usize) -> Self {
        use std::alloc::{alloc, Layout};
        
        let layout = unsafe { Layout::from_size_align_unchecked(capacity, 1) };
        let data = unsafe { alloc(layout) };

        Self {
            data: NonNull::new(data).unwrap(),
            cap: capacity,
            locations: Vec::new(),
        }
    }

    /// Clears the `Allocal`, invalidating all locations.
    pub fn clear(&mut self) {
        for loc in self.locations.drain(..) {
            if let Some(drop_fn) = loc.drop_fn {
                // The object needs to be dropped
                let ptr = unsafe { self.data.as_ptr().add(loc.offset) };
                drop_fn(ptr);
            }
        }
    }

    /// Reserves some memory to allocate `count` more distincts items into the
    /// `Allocal`. Note that this does not increese the capacity of the `Allocal`:
    /// this function basically calls `Vec::reserve` on the vector that keeps
    /// track of the locations of this `Allocal`.
    pub fn reserve_items(&mut self, count: usize) {
        self.locations.reserve(count);
    }

    /// Gets a `*const T` to a `T` stored in the `Allocal`.
    #[inline]
    pub fn get_ptr<T>(&self, location: Location<T>) -> Option<*const T> {
        self.locations.binary_search_by_key(&location.offset, |loc| loc.offset).ok()?;
        
        unsafe {
            Some(self.data.as_ptr().add(location.offset).cast())
        }
    }

    /// Gets a `*mut T` to a `T` stored in the `Allocal`.
    #[inline]
    pub fn get_mut_ptr<T>(&mut self, location: Location<T>) -> Option<*mut T> {
        self.locations.binary_search_by_key(&location.offset, |loc| loc.offset).ok()?;
        
        unsafe {
            Some(self.data.as_ptr().add(location.offset).cast())
        }
    }

    /// Gets a reference to an item stored in the `Allocal`.
    #[inline(always)]
    pub fn try_get<T>(&self, location: Location<T>) -> Option<&T> {
        unsafe {
            Some(&*self.get_ptr(location)?)
        }
    }
    
    /// Gets a mutable reference to an item stored in the `Allocal`.
    #[inline(always)]
    pub fn try_get_mut<T>(&mut self, location: Location<T>) -> Option<&mut T> {
        unsafe {
            Some(&mut *self.get_mut_ptr(location)?)
        }
    }

    /// Gets a reference to an item stored in the `Allocal`.
    #[inline(always)]
    pub fn get<T>(&self, location: Location<T>) -> &T {
        self.try_get(location).unwrap()
    }
    
    /// Gets a mutable reference to an item stored in the `Allocal`.
    #[inline(always)]
    pub fn get_mut<T>(&mut self, location: Location<T>) -> &mut T {
        self.try_get_mut(location).unwrap()
    }

    /// Allocates a new `T` in this `Allocal` and returns its location.
    /// If the `Allocal` is full, `Err(())` is returned.
    pub fn allocate<T: 'static>(&mut self, item: T) -> Result<Location<T>, ()> {
        use std::mem::{align_of, size_of, needs_drop};

        let size = size_of::<T>();

        // Handle zero-sized structures.
        if size == 0 {
            let location = LocationInner {
                size,
                offset: usize::MAX,
                drop_fn: None,
            };

            self.locations.push(location);

            // Nothing needs to be written as this is a zero sized type.
            
            return Ok(Location {
                offset: usize::MAX,
                _m: PhantomData,
            });
        }

        let align = align_of::<T>();

        // (self.data + offset) is aligned !
        let mut offset = self.data.as_ptr() as usize % align;

        let loc = if self.locations.is_empty() {
            // The `Localloc` is currently empty.
            // We just need to check whether it is large enough to store our `item`.
            if offset + size > self.cap {
                // Our allocated data is too small.
                return Err(());
            }
            
            0usize
        } else {
            let mut loc = 0usize;

            loop {
                while offset < self.locations[loc].offset + self.locations[loc].size {
                    offset += align;
                }

                // (self.data + offset) is still aligned

                loc += 1;

                if loc >= self.locations.len() {
                    // We are at the end of the allocated data.
                    // We just need to check whether our remaining allocated data
                    // is large enough.
                    if offset + size > self.cap {
                        // Too small!
                        return Err(());
                    }

                    break;
                }

                if offset + size < self.locations[loc].offset {
                    // We found an aligned location between two other taken locations.
                    break;
                }
            }

            loc
        };

        // We can insert our location
        let location = LocationInner {
            offset,
            size,
            drop_fn: if needs_drop::<T>() {
                Some(|ptr: *mut u8| unsafe { std::ptr::drop_in_place(ptr as *mut T) })
            } else { None },
        };

        self.locations.insert(loc, location.clone());

        // We can safely write our item because the found location is known
        // to be uninitialized.
        unsafe {
            let ptr: *mut T = self.data.as_ptr().add(offset).cast();
            ptr.write(item);
        }

        Ok(Location {
            offset,
            _m: PhantomData,
        })
    }

    /// Deallocate an existing object stored at the given location.
    /// 
    /// # Panics
    /// This function panics if the given location is invalid.
    pub fn deallocate<T>(&mut self, location: Location<T>) -> T {
        let idx = match self.locations.binary_search_by_key(&location.offset, |loc| loc.offset) {
            Ok(val) => val,
            Err(_) => panic!("Invalid location"),
        };
        
        let location = self.locations.remove(idx);

        // Dropping the object is up the caller
        unsafe {
            let ptr: *mut T = self.data.as_ptr().add(location.offset).cast();
            ptr.read()
        }
    }
}

#[cfg(test)]
#[test]
fn create() {
    let mut allocal = Allocal::with_capacity(100);
    let mut locations = Vec::new();

    for i in 0u8..100 {
        locations.push(allocal.allocate(i).unwrap());
    }

    assert!(allocal.allocate(0u8).is_err());

    for (i, loc) in locations.clone().into_iter().enumerate() {
        assert_eq!(*allocal.get(loc), i as u8);
    }

    for (i, loc) in locations.clone().into_iter().enumerate().rev() {
        assert_eq!(allocal.deallocate(loc), i as u8);
    }

    for loc in locations.into_iter() {
        assert_eq!(allocal.try_get(loc), None);
    }
}

#[cfg(test)]
#[test]
fn alignement() {
    use std::mem::align_of;

    for _ in 0..100 {
        let mut allocal = Allocal::with_capacity(30);
        let loc = allocal.allocate(0u128).unwrap();
        
        let ptr = allocal.get_ptr(loc).unwrap() as usize;
        assert_eq!(ptr % align_of::<u128>(), 0);
    }
}

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

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

    {
        let mut allocal = Allocal::with_capacity(1024);

        for _ in 0..20 {
            allocal.allocate(DropCheck(Rc::clone(&drop_count))).unwrap();
        }
    }

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