meowalloc 0.1.1

use meowvec::MeowVec;

use core::{alloc::Layout, cmp::Ordering, num::NonZero, ptr::NonNull};

use crate::{block::{Block, BlockError, BlockIdxGuard, ExtractIdx}};

pub struct BlockStore<const MAX_BLOCKS: usize>(pub MeowVec<Block, MAX_BLOCKS>);

impl<const MAX_BLOCKS: usize> BlockStore<MAX_BLOCKS> {
    pub fn new(init_ptr: NonNull<u8>, init_size: NonZero<usize>) -> Self {
        let mut svec = MeowVec::new();
        svec.push(Block { ptr: init_ptr, size: init_size });
        Self(svec)
    }
    /// Splits a block in two, leaving the part starting at current ptr at the current block_idx
    /// and pushing the part starting at offset to the block list
    ///
    /// If the offset is zero or equal to the size of the block, no new block is created.
    /// 
    /// Returns two BlockIdxGuards of the split block(s) on success
    pub fn split(&mut self, block_idx: usize, offset: usize) -> Result<(BlockIdxGuard<'_>, BlockIdxGuard<'_>), BlockError> {
        let left = self.0.get_mut(block_idx)
            .ok_or(BlockError::InvalidIndex)?;
        let r_size = left.size.get().checked_sub(offset)
            .ok_or(BlockError::TooSmall)?;
        if r_size == 0 || offset == 0 {
            return Ok((BlockIdxGuard::new(block_idx), BlockIdxGuard::new(block_idx)))
        }
        let (r_size, offset) = unsafe {(
            NonZero::new_unchecked(r_size),
            NonZero::new_unchecked(offset)
        )};
        left.size = offset;
        
        let ptr = unsafe { left.ptr.add(offset.get()) };
        let right = Block { ptr, size: r_size };

        self.0.try_push(right)
            .map_err(|_| BlockError::Overflow)?;
        let blocks = &self.0;
        let right_idx = blocks.len() - 1;

        Ok((BlockIdxGuard::new(block_idx), BlockIdxGuard::new(right_idx)))
    }

    /// Merges a block with another adjacent block
    pub fn merge(&mut self, block_idx_1: usize, block_idx_2: usize) -> Result<BlockIdxGuard<'_>, BlockError> {
        let (b1, b2) = (
            self.0.get(block_idx_1).ok_or(BlockError::InvalidIndex)?,
            self.0.get(block_idx_2).ok_or(BlockError::InvalidIndex)?
        );
        #[cfg(test)]
        dbg!(b1, b2);
        let idx = match b1.cmp_adjacent(b2).ok_or(BlockError::NotAdjacent)? {
            Ordering::Less => unsafe { self.merge_inner(block_idx_1, block_idx_2) },
            Ordering::Greater => unsafe { self.merge_inner(block_idx_2, block_idx_1) },
            Ordering::Equal => unreachable!()
        };
        
        Ok(BlockIdxGuard::new(idx))
    }

    /// Merges the lesser-adjacent block with the greater adjacent block and returns the index of the
    /// merged block
    ///
    /// # Safety:
    /// The caller must ensure that:
    /// - Block indices refer to real blocks in the store
    /// - Block under l_block_idx is lesser-adjacent
    /// - The resulting index is not used after being invalidated
    /// - Usize overflow cannot occur when adding block sizes
    unsafe fn merge_inner(&mut self, l_block_idx: usize, g_block_idx: usize) -> usize {
        unsafe {
            let greater = self.0.get_unchecked(g_block_idx);
            let g_size = greater.size.get();
            let lesser = self.0.get_unchecked_mut(l_block_idx);
            lesser.size = lesser.size.unchecked_add(g_size);
            self.0.swap_remove_unchecked(g_block_idx);
            // l_block was last, swap index
            let l_block_idx = if l_block_idx == self.0.len() { g_block_idx }
            else { l_block_idx };
            l_block_idx
        }
    }

    /// Splits off a block from a suitably sized block from the selection, with the specified layout
    /// # Safety:
    /// The caller must ensure that:
    /// - The returned index is not used after the block store
    /// has been mutated in such a way as to make it point to another block entirely
    pub unsafe fn create(&mut self, layout: Layout, selection: impl IntoIterator<Item = usize>) -> Result<usize, BlockError> {
        let layout = layout.pad_to_align();
        
        for idx in selection {
            let block = self.0.get(idx).ok_or(BlockError::InvalidIndex)?;
            
            let (raw, _) = block.size.get().overflowing_sub(layout.size());
            let offset = raw - raw % layout.align();
            
            match self.split(idx, offset) {
                Ok(guard) => return unsafe { Ok(guard.1.extract()) },
                Err(BlockError::TooSmall) => (),
                Err(e) => return Err(e)
            };
        }
        
        Err(BlockError::TooSmall)
    }

    /// Returns a `BlockIdxGuard` of the block from the selection whose `block.ptr` matches `ptr` if it exists
    ///
    /// Short-circuits on the first index out of range
    #[inline(always)]
    pub fn find_ptr(&self, ptr: NonNull<u8>, selection: impl IntoIterator<Item = usize>) -> Option<BlockIdxGuard<'_>> {
        for idx in selection {
            let block = self.0.get(idx)?;
            if block.ptr == ptr {
                return Some(BlockIdxGuard::new(idx))
            }
        }
        None
    }

    /// Returns an iterator of `BlockIdxGuard`s immediately adjacent to `block`
    ///
    /// Short-circuits on the first index out of range
    #[inline(always)]
    pub fn find_adjacent(&self, block: &Block, selection: impl IntoIterator<Item = usize>) -> impl Iterator<Item = BlockIdxGuard<'_>> {
        selection
            .into_iter()
            .take_while(|&idx| idx < self.len())
            .map(|idx| (idx, unsafe { self.0.get_unchecked(idx) } ))
            .filter(|(_, mb_adjacent)| mb_adjacent.is_adjacent(block))
            .map(|(idx, _)| BlockIdxGuard::new(idx))
    }

    pub const fn len(&self) -> usize { self.0.len() }
}

#[cfg(test)]
extern crate std;
#[cfg(test)]
use std::dbg;

#[test]
fn find_ptr() {
    let ptr = NonNull::new(1 as *mut u8).unwrap();
    let size = NonZero::new(1).unwrap();
    let mut blocks = MeowVec::<Block, 8>::new();
    blocks.push(Block { ptr: unsafe { ptr.add(2) }, size });
    blocks.push(Block { ptr, size });
    blocks.push(Block { ptr: unsafe { ptr.add(1) }, size });
    let store = BlockStore(blocks);
    
    assert_eq!(
        unsafe { store.find_ptr(ptr, [0, 1, 2]).unwrap().extract() },
        1
    );
    assert_eq!(
        unsafe { store.find_ptr(ptr.add(1), [0, 1, 2]).unwrap().extract() },
        2
    );
    assert_eq!(
        unsafe { store.find_ptr(ptr.add(2), [0, 1, 2]).unwrap().extract() },
        0
    );
    assert!(store.find_ptr(ptr, [0, 2]).is_none());
}

#[test]
fn find_adjacent() {
    let ptr = NonNull::new(1 as *mut u8).unwrap();
    let mut blocks = MeowVec::<Block, 8>::new();
    blocks.push(Block { ptr: unsafe { ptr.add(2) }, size: NonZero::new(2).unwrap() });
    blocks.push(Block { ptr, size: NonZero::new(2).unwrap() });
    blocks.push(Block { ptr: unsafe { ptr.add(4) }, size: NonZero::new(4).unwrap() });
    blocks.push(Block { ptr: unsafe { ptr.add(8) }, size: NonZero::new(1).unwrap() });
    let store = BlockStore(blocks);
    
    let indices = unsafe {
        store.find_adjacent(&store.0[0], [0, 1, 2, 3])
            .map(|guard| guard.extract())
            .collect::<MeowVec<usize, 8>>()
    };
    assert_eq!(indices, (&[1, 2][..]).into());
    
    assert!(unsafe { store.find_adjacent(&store.0[0], [0, 3]) }.next().is_none());

    let indices = unsafe {
        store.find_adjacent(&store.0[2], [0, 1, 2, 3])
            .map(|guard| guard.extract())
            .collect::<MeowVec<usize, 8>>()
    };

    assert_eq!(indices, (&[0, 3][..]).into());
}

#[test]
fn split() {
    let ptr = NonNull::new(1 as *mut u8).unwrap();
    let mut store = BlockStore::<3>::new(ptr, NonZero::new(64).unwrap());
    
    let (parent, child) = unsafe { store.split(0, 8).unwrap().extract() };
    let init_block = parent;
    let (parent, child) = (&store.0[parent], &store.0[child]);
    assert_eq!(parent.ptr, ptr);
    assert_eq!(parent.size, NonZero::new(8).unwrap());
    assert_eq!(child.ptr, unsafe { ptr.add(8) });
    assert_eq!(child.size, NonZero::new(56).unwrap());

    let (parent, child) = unsafe { store.split(1, 16).unwrap().extract() };
    let (parent, child) = (&store.0[parent], &store.0[child]);
    let init_block = &store.0[init_block];
    
    // Didn't change old block
    assert_eq!(init_block.ptr, ptr);
    assert_eq!(init_block.size, NonZero::new(8).unwrap());

    assert_eq!(parent.ptr, unsafe { ptr.add(8) });
    assert_eq!(parent.size, NonZero::new(16).unwrap());
    assert_eq!(child.ptr, unsafe { ptr.add(24) });
    assert_eq!(child.size, NonZero::new(40).unwrap());

    assert_eq!(store.split(3, 1).unwrap_err(), BlockError::InvalidIndex);
    assert_eq!(store.split(2, 41).unwrap_err(), BlockError::TooSmall);
    assert_eq!(store.split(0, 2).unwrap_err(), BlockError::Overflow);
}

#[test]
fn merge() {
    let ptr = NonNull::new(1 as *mut u8).unwrap();
    let mut blocks = MeowVec::<Block, 8>::new();
    blocks.push(Block { ptr: unsafe { ptr.add(2) }, size: NonZero::new(2).unwrap() });
    blocks.push(Block { ptr, size: NonZero::new(2).unwrap() });
    blocks.push(Block { ptr: unsafe { ptr.add(4) }, size: NonZero::new(4).unwrap() });
    blocks.push(Block { ptr: unsafe { ptr.add(8) }, size: NonZero::new(2).unwrap() });
    let mut store = BlockStore(blocks);
     
    assert_eq!(store.merge(3, 4).unwrap_err(), BlockError::InvalidIndex);
    assert_eq!(store.merge(1, 2).unwrap_err(), BlockError::NotAdjacent);
    
    let idx_1 = unsafe { store.merge(0, 1).unwrap().extract() };
    assert_eq!(store.0.len(), 3);
    assert_eq!(store.0[idx_1].ptr, ptr);
    assert_eq!(store.0[idx_1].size, NonZero::new(4).unwrap());
    assert_eq!(idx_1, 1);
    
    // blocks[3] is now at index 0 thanks to swapping positions with greater-adjacent blocks[0] (now nonexistent)
    assert_eq!(store.0[0].ptr, unsafe { ptr.add(8) });
    assert_eq!(store.0[0].size, NonZero::new(2).unwrap());
    
    let idx_2 = unsafe { store.merge(2, 0).unwrap().extract() };
    assert_eq!(store.0.len(), 2);
    assert_eq!(store.0[idx_2].ptr, unsafe { ptr.add(4) });
    assert_eq!(store.0[idx_2].size, NonZero::new(6).unwrap());
    assert_eq!(idx_2, 0);

    let idx = unsafe { store.merge(idx_1, idx_2).unwrap().extract() };
    assert_eq!(store.0.len(), 1);
    assert_eq!(store.0[idx_2].ptr, ptr);
    assert_eq!(store.0[idx_2].size, NonZero::new(10).unwrap());
    assert_eq!(idx, 0);
}

#[test]
fn create() {
    let init_ptr = NonNull::new(1 as *mut u8).unwrap();
    let init_size = NonZero::new(256).unwrap();
    
    let align_1 = Layout::new::<[u8; 8]>();
    let align_2 = Layout::new::<[u16; 4]>();
    let align_8 = Layout::new::<[u64; 8]>();
    let align_16 = Layout::new::<[u128; 8]>();
    let fat_ass = Layout::new::<[u128; 1024]>();

    let expected_sizes = [
        size_of::<[u8; 8]>(),
        size_of::<[u16; 4]>(),
        size_of::<[u64; 8]>(),
        size_of::<[u128; 8]>(),
    ];

    let expected_sizes = expected_sizes.map(|n| NonZero::new(n).unwrap());

    let mut store = BlockStore::<8>::new(init_ptr, init_size);
    
    unsafe {
        store.create(align_1, [0]).unwrap();
        store.create(align_2, [0]).unwrap();
        store.create(align_8, [0]).unwrap();
        store.create(align_16, [0]).unwrap();
        store.create(fat_ass, [0]).unwrap_err();
    };

    let pointers: MeowVec<NonNull<u8>, 5> = store.0[2..].into_iter().map(|block| block.ptr).collect();
    let sizes: MeowVec<NonZero<usize>, 5> = store.0[2..].into_iter().map(|block| block.size).collect();

    let size = store.0[1].size.get();
    assert!(size >= expected_sizes[0].get());
    let expected_ptr = unsafe { init_ptr.add(init_size.get()).sub(size) };
    assert_eq!(store.0[1].ptr, expected_ptr);
    
    let mut prev_ptr = store.0[1].ptr;
    
    for ((size, expected), ptr) in sizes.iter().zip(expected_sizes.iter()).zip(pointers.iter()) {
        assert!(size >= expected);
        assert_eq!(unsafe { prev_ptr.sub(size.get()) }, *ptr);
        prev_ptr = *ptr;
    }
}