yggdrasil/lib/libyalloc/src/allocator.rs

use core::{alloc::Layout, ops::Index, ptr::NonNull};

use crate::{
    bucket::Bucket,
    sys::{self, PageProvider},
    util::{Assert, IsTrue},
};

struct BucketList<P: PageProvider, const N: usize, const M: usize>
where
    [u64; M / 64]: Sized,
    Assert<{ M % 64 == 0 }>: IsTrue,
{
    head: Option<NonNull<Bucket<P, N, M>>>,
}

pub struct BucketAllocator<P: PageProvider> {
    // 1024x64 = 16 pages
    buckets_1024: BucketList<P, 1024, 64>,
    // 512x64 = 8 pages
    buckets_512: BucketList<P, 512, 64>,
    // 256x128 = 8 pages
    buckets_256: BucketList<P, 256, 128>,
    // 128x128 = 4 pages
    buckets_128: BucketList<P, 128, 128>,
    // 64x128 = 2 pages
    buckets_64: BucketList<P, 128, 128>,
    // 32x256 = 2 pages
    buckets_32: BucketList<P, 128, 128>,
}

impl<P: PageProvider, const N: usize, const M: usize> BucketList<P, N, M>
where
    [u64; M / 64]: Sized,
    Assert<{ M % 64 == 0 }>: IsTrue,
{
    const fn new() -> Self {
        Self { head: None }
    }

    fn allocate(&mut self) -> Option<NonNull<u8>> {
        let mut node = self.head;
        while let Some(mut bucket) = node {
            let bucket = unsafe { bucket.as_mut() };

            if let Some(ptr) = bucket.allocate() {
                return Some(ptr);
            }

            node = bucket.next;
        }

        // No usable bucket found
        let mut node = Bucket::new()?;
        let bucket = unsafe { node.as_mut() };
        bucket.next = self.head;
        self.head = Some(node);

        bucket.allocate()
    }

    unsafe fn free(&mut self, ptr: NonNull<u8>) {
        let mut node = self.head;
        while let Some(mut bucket) = node {
            let bucket = bucket.as_mut();

            if let (true, _last) = bucket.free(ptr) {
                // TODO free the node if last?
                return;
            }

            node = bucket.next;
        }

        panic!("Possible double free detected: pointer {:p} from bucket list {}x{}B, no corresponding bucket found", ptr, N, M);
    }
}

impl<P: PageProvider, const N: usize, const M: usize> Index<usize> for BucketList<P, N, M>
where
    [u64; M / 64]: Sized,
    Assert<{ M % 64 == 0 }>: IsTrue,
{
    type Output = Bucket<P, N, M>;

    fn index(&self, index: usize) -> &Self::Output {
        let mut current = 0;
        let mut node = self.head;
        while let Some(bucket) = node {
            let bucket = unsafe { bucket.as_ref() };
            if current == index {
                return bucket;
            }
            current += 1;
            node = bucket.next;
        }

        panic!(
            "BucketList index out of range: contains {} buckets, tried to index {}",
            current, index
        );
    }
}

impl<P: PageProvider> BucketAllocator<P> {
    pub const fn new() -> Self {
        Self {
            buckets_1024: BucketList::new(),
            buckets_512: BucketList::new(),
            buckets_256: BucketList::new(),
            buckets_128: BucketList::new(),
            buckets_64: BucketList::new(),
            buckets_32: BucketList::new(),
        }
    }

    pub fn allocate(&mut self, layout: Layout) -> Option<NonNull<u8>> {
        let aligned = layout.pad_to_align();

        match aligned.size() {
            size if size <= 32 => self.buckets_32.allocate(),
            size if size <= 64 => self.buckets_64.allocate(),
            size if size <= 128 => self.buckets_128.allocate(),
            size if size <= 256 => self.buckets_256.allocate(),
            size if size <= 512 => self.buckets_512.allocate(),
            size if size <= 1024 => self.buckets_1024.allocate(),
            size => P::map_pages(size.div_ceil(sys::PAGE_SIZE)),
        }
    }

    /// # Safety
    ///
    /// Unsafe: may cause double-free problem and panic
    pub unsafe fn free(&mut self, ptr: NonNull<u8>, layout: Layout) {
        let aligned = layout.pad_to_align();

        match aligned.size() {
            size if size <= 32 => self.buckets_32.free(ptr),
            size if size <= 64 => self.buckets_64.free(ptr),
            size if size <= 128 => self.buckets_128.free(ptr),
            size if size <= 256 => self.buckets_256.free(ptr),
            size if size <= 512 => self.buckets_512.free(ptr),
            size if size <= 1024 => self.buckets_1024.free(ptr),
            size => {
                assert_eq!(usize::from(ptr.addr()) % sys::PAGE_SIZE, 0);
                P::unmap_pages(ptr, size.div_ceil(sys::PAGE_SIZE));
            }
        }
    }
}

unsafe impl<P: PageProvider> Sync for BucketAllocator<P> {}

#[cfg(test)]
mod tests {
    use core::{alloc::Layout, ptr::NonNull};

    use super::{BucketAllocator, BucketList};
    use crate::sys::OsPageProvider;

    #[test]
    fn single_list_allocation() {
        let mut list = BucketList::<OsPageProvider, 32, 64>::new();
        let mut vec = vec![];

        for _ in 0..4 * 64 + 3 {
            let ptr = list.allocate().unwrap();
            vec.push(ptr);
        }

        for ptr in vec {
            unsafe {
                list.free(ptr);
            }
        }
    }

    #[test]
    fn multi_list_allocation() {
        const SIZES: &[usize] = &[1, 3, 7, 15, 16, 24, 33, 65, 126, 255, 500, 1000];

        let mut allocator = BucketAllocator::<OsPageProvider>::new();
        let mut vec = vec![];

        for _ in 0..65 {
            for &size in SIZES {
                let layout = Layout::from_size_align(size, 16).unwrap();
                let ptr = allocator.allocate(layout).unwrap();
                assert_eq!(usize::from(ptr.addr()) % 16, 0);
                let mut slice = NonNull::slice_from_raw_parts(ptr, size);
                unsafe {
                    slice.as_mut().fill(123);
                }
                vec.push((ptr, layout));
            }
        }

        for (ptr, layout) in vec {
            unsafe {
                allocator.free(ptr, layout);
            }
        }
    }

    #[test]
    #[should_panic]
    fn double_free() {
        let mut allocator = BucketAllocator::<OsPageProvider>::new();
        let layout = Layout::from_size_align(63, 32).unwrap();
        let ptr = allocator.allocate(layout).unwrap();

        unsafe {
            allocator.free(ptr, layout);
            allocator.free(ptr, layout);
        }
    }

    #[test]
    fn large_alloc() {
        const SIZES: &[usize] = &[2000, 2048, 4000, 4096, 8192];

        let mut allocator = BucketAllocator::<OsPageProvider>::new();
        let mut vec = vec![];

        for &size in SIZES {
            let layout = Layout::from_size_align(size, 32).unwrap();
            let ptr = allocator.allocate(layout).unwrap();
            vec.push((ptr, layout));
        }

        for (ptr, layout) in vec {
            assert_eq!(usize::from(ptr.addr()) % 0x1000, 0);

            unsafe {
                allocator.free(ptr, layout);
            }
        }
    }
}