meowalloc 0.1.0

use kitset::KitSet;

use core::{alloc::{AllocError, Layout, Allocator, GlobalAllocator}, num::NonZero, ptr::NonNull};

#[cfg(test)] use crate::meow_allocator;
use crate::{MeowAllocator128, atomic_table::AtomicTable, constants::linux::*, spinlock::SpinLock, superblock::Superblock};

use crate as meowalloc;

const MIN_SUPERBLOCK_PAGES: NonZero<usize> = NonZero::new(64).unwrap();
const MIN_SUPERBLOCK_SIZE: NonZero<usize> = PAGE_SIZE.checked_mul(MIN_SUPERBLOCK_PAGES).unwrap();
const TRIES_BEFORE_NEW_SUPERBLOCK: usize = 16;
const TRIES_BEFORE_SPIN_DEALLOC: usize = 16;


pub struct MeowAllocator<
    const MAX_SUPERBLOCKS: usize,
    const TABLE_BITSET_WORDS: usize,
    const MAX_BLOCKS_PER_SUPER: usize,
    const FREE_BITSET_WORDS: usize,
> {
    superblocks: AtomicTable<
        SpinLock<Superblock<MAX_BLOCKS_PER_SUPER, FREE_BITSET_WORDS>>,
        MAX_SUPERBLOCKS, TABLE_BITSET_WORDS
    >
}

impl<
    const MAX_SUPERBLOCKS: usize,
    const TABLE_BITSET_WORDS: usize,
    const MAX_BLOCKS_PER_SUPER: usize,
    const FREE_BITSET_WORDS: usize,
> MeowAllocator<MAX_SUPERBLOCKS, TABLE_BITSET_WORDS, MAX_BLOCKS_PER_SUPER, FREE_BITSET_WORDS> {
    fn alloc_new_superblock(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        #[cfg(test)]
        dbg!(size_of::<Self>());
        let sblock_size = layout.size().max(MIN_SUPERBLOCK_SIZE.get());
        let sblock_size = unsafe { NonZero::new_unchecked(sblock_size) };
        let mut sblock = Superblock::new(sblock_size)?;
        let ptr = sblock.alloc(layout)?;
        self.superblocks.add(SpinLock::new(sblock))
            .map_err(|_| AllocError)?;
        Ok(ptr)
    }

    /// Attempts to allocate a layout on one of the currently unlocked superblocks
    /// if they aren't in the visited set.
    /// Adds visited unlocked blocks to the set
    fn try_alloc(&self, layout: Layout, visited: &mut KitSet<TABLE_BITSET_WORDS>) -> Result<NonNull<[u8]>, AllocError> {
        for (idx, sblock) in self.superblocks.iter() {
            if visited.is_one(idx) { continue }
            
            if let Some(mut handle) = sblock.try_lock() {
                match handle.alloc(layout) {
                    Ok(ptr) => return Ok(ptr),
                    Err(_) => visited.set_one(idx)
                }
            }
        }
        
        Err(AllocError)
    }

    /// Attempts to deallocate a layout at ptr on one of the currently unlocked superblocks
    /// if they aren't in the visited set.
    /// Adds visited blocks to the set
    unsafe fn try_dealloc(&self, ptr: NonNull<u8>, _: Layout, visited: &mut KitSet<TABLE_BITSET_WORDS>) -> Result<(), AllocError> {
        for (idx, superblock) in self.superblocks.iter() {
            if visited.is_one(idx) { continue }
            
            if let Some(mut handle) = superblock.try_lock() {
                match unsafe { handle.dealloc(ptr) } {
                    Ok(_) => return Ok(()),
                    Err(_) => visited.set_one(idx)
                }
            }
        }
        Err(AllocError)
    }

    /// Attempts to deallocate layout at ptr in a best effort manner
    ///
    /// Iterates over all the unvisited superblocks, spinning until the lock is acquired on every
    /// iteration, until a successful deallocation has happened or there are no more unvisited superblocks.
    unsafe fn spin_dealloc(&self, ptr: NonNull<u8>, _: Layout, visited: KitSet<TABLE_BITSET_WORDS>) {
        let iter = self.superblocks
            .iter()
            .filter_map(|(idx, sblock)| visited.is_zero(idx).then_some(sblock));

        for superblock in iter {
            if unsafe { superblock.lock().dealloc(ptr).is_ok() } {
                return
            }
        }
    }
}

unsafe impl<
    const MAX_SUPERBLOCKS: usize,
    const TABLE_BITSET_WORDS: usize,
    const MAX_BLOCKS_PER_SUPER: usize,
    const FREE_BITSET_WORDS: usize,
> Allocator for MeowAllocator<MAX_SUPERBLOCKS, TABLE_BITSET_WORDS, MAX_BLOCKS_PER_SUPER, FREE_BITSET_WORDS> {
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        if layout.alignment() > PAGE_ALIGNMENT {
            return Err(AllocError)
        }
        let mut visited = KitSet::new();
        // Search for an unlocked superblock that has enough space for allocation
        for _ in 0..TRIES_BEFORE_NEW_SUPERBLOCK {
            match self.try_alloc(layout, &mut visited) {
                Ok(ptr) => return Ok(ptr),
                Err(_) => continue
            }
        };

        // Add one more superblock in case of high lock contention or not enough memory
        self.alloc_new_superblock(layout)
    }
   
    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
        let mut visited = KitSet::new();
       
        // Search for an unlocked superblock that has the allocation
        for _ in 0..TRIES_BEFORE_SPIN_DEALLOC {
            match unsafe { self.try_dealloc(ptr, layout, &mut visited) } {
                Ok(_) => return,
                Err(_) => continue
            }
        }

        // Pick up leftovers (in case of high lock contention)
        unsafe { self.spin_dealloc(ptr, layout, visited) };
    }
}

unsafe impl<
    const MAX_SUPERBLOCKS: usize,
    const TABLE_BITSET_WORDS: usize,
    const MAX_BLOCKS_PER_SUPER: usize,
    const FREE_BITSET_WORDS: usize,
> GlobalAllocator for MeowAllocator<MAX_SUPERBLOCKS, TABLE_BITSET_WORDS, MAX_BLOCKS_PER_SUPER, FREE_BITSET_WORDS> {}

impl<
    const MAX_SUPERBLOCKS: usize,
    const TABLE_BITSET_WORDS: usize,
    const MAX_BLOCKS_PER_SUPER: usize,
    const FREE_BITSET_WORDS: usize,
> MeowAllocator<MAX_SUPERBLOCKS, TABLE_BITSET_WORDS, MAX_BLOCKS_PER_SUPER, FREE_BITSET_WORDS> {
    pub fn new() -> Self {
        Self { superblocks: AtomicTable::new() }
    }
}

#[cfg(test)]
extern crate std;
#[cfg(test)]
use std::{dbg, rc::Rc};

#[cfg(test)]
const LAYOUT_2KIB_16: Layout = Layout::new::<[u128; 128]>();
#[cfg(test)]
const LAYOUT_1MIB_1: Layout = Layout::new::<[u8; 1_048_576]>();
#[cfg(test)]
const LAYOUT_ZST: Layout = Layout::new::<()>();

#[cfg(test)]
type MeowAllocator16 = meow_allocator!(16usize, 16usize);

#[test]
fn alloc_new_superblock_1mib_1() {
    let allocator = Rc::new(MeowAllocator16::new());
    let ptr = allocator.alloc_new_superblock(LAYOUT_1MIB_1).unwrap();
    let sblock = allocator.superblocks.get(0).unwrap();
    let handle = sblock.lock();
    let block = &handle.blocks().0[0];
    assert_eq!(block.size.get(), LAYOUT_1MIB_1.size());
    assert_eq!(block.ptr, ptr.cast::<u8>());
}

#[test]
fn alloc_new_superblock_2kib_16() {
    let allocator = Rc::new(MeowAllocator16::new());
    let ptr = allocator.alloc_new_superblock(LAYOUT_2KIB_16).unwrap();
    let sblock = allocator.superblocks.get(0).unwrap();
    let handle = sblock.lock();
    let block = &handle.blocks().0[1];
    assert_eq!(block.size.get(), LAYOUT_2KIB_16.size());
    assert_eq!(block.ptr, ptr.cast::<u8>());
}

#[test]
fn alloc_new_superblock_zst() {
    let allocator = Rc::new(MeowAllocator16::new());
    allocator.alloc_new_superblock(LAYOUT_ZST).unwrap();
    let sblock = allocator.superblocks.get(0).unwrap();
    let handle = sblock.lock();
    let block = &handle.blocks().0[0];
    assert_eq!(block.size, MIN_SUPERBLOCK_SIZE);
}

#[test]
fn try_alloc_uncontended() {
    let allocator = Rc::new(MeowAllocator16::new());
    let mut visited = KitSet::zeros();
   
    // No superblocks, fails
    allocator.try_alloc(LAYOUT_2KIB_16, &mut visited).unwrap_err();
   
    // Too large, fails
    allocator.alloc_new_superblock(LAYOUT_ZST).unwrap();
    allocator.try_alloc(LAYOUT_1MIB_1, &mut visited).unwrap_err();
    assert!(visited.is_one(0));
   
    // Only superblock marked visited, fails
    allocator.try_alloc(LAYOUT_2KIB_16, &mut visited).unwrap_err();
    assert!(visited.is_one(0));
   
    // Only superblock is unlocked and is large enough, succeeds
    visited.flip(0);
    allocator.try_alloc(LAYOUT_2KIB_16, &mut visited).unwrap();
    assert!(visited.is_zero(0));
   
    // large enough and unlocked superblock exists, succeeds
    allocator.alloc_new_superblock(LAYOUT_ZST).unwrap();
    visited.flip(0);
    allocator.try_alloc(LAYOUT_2KIB_16, &mut visited).unwrap();
    assert!(visited.is_one(0));
    assert!(visited.is_zero(1));
    visited.flip(0)
}

#[test]
fn try_dealloc() {
    let allocator = Rc::new(MeowAllocator16::new());
    let mut visited_dealloc = KitSet::zeros();
    let dangling = LAYOUT_1MIB_1.dangling_ptr();

    allocator.alloc_new_superblock(LAYOUT_1MIB_1).unwrap();
    let ptr = allocator.alloc_new_superblock(LAYOUT_1MIB_1).unwrap().cast::<u8>();
    
    // Alien pointer, fails
    let res = unsafe { allocator.try_dealloc(dangling, LAYOUT_1MIB_1, &mut visited_dealloc) };
    dbg!(&res);
    assert_eq!(res, Err(AllocError));
    assert!(visited_dealloc.is_one(0));
    assert!(visited_dealloc.is_one(1));

    // Superblock containing ptr visited, fails 
    unsafe { allocator.try_dealloc(ptr, LAYOUT_1MIB_1, &mut visited_dealloc).unwrap_err() };
    
    // Superblock containing ptr unlocked and unvisited, succeeds
    visited_dealloc.flip(0);
    visited_dealloc.flip(1);
    assert!(unsafe { allocator.try_dealloc(ptr, LAYOUT_1MIB_1, &mut visited_dealloc).is_ok() });
    assert!(visited_dealloc.is_one(0));
    assert!(visited_dealloc.is_zero(1))
}