alloc_cat 1.1.1

use core::fmt;
use core::ptr;

use crate::SimpleAllocatorStats;


// max size of a block that can be tracked in the free list. for larger
// allocations, allocate 64KB pages directly.
pub const MAX_ALLOC: usize = 32 * 1024 - 1;

// max distance that can be tracked between blocks in the free list. for a
// larger heap, chain multiple SimpleAlloc's together.
pub const MAX_HEAP: usize = 2 * 1024 * 1024;


// crazy packed info for each node in a free block (must fit into one word)
#[derive(Debug)]
pub struct FreeLink {
    // size in words (high 13 bits = 32KB), word-offset to next (low 19 bits = 2MB)
    packed: u32,
}

impl FreeLink {
    fn bless(ptr: *mut u32) -> &'static mut FreeLink {
        unsafe { &mut *(ptr as *mut FreeLink) }
    }

    fn jump(&self, offset: isize) -> &'static FreeLink {
        unsafe { &*((self as *const FreeLink).offset(offset)) }
    }

    fn jump_mut(&mut self, offset: isize) -> &'static mut FreeLink {
        unsafe { &mut *((self as *mut FreeLink).offset(offset)) }
    }

    fn compute_jump_to(&self, other: &FreeLink) -> isize {
        unsafe { (other as *const FreeLink).offset_from(self as *const FreeLink) }
    }

    fn init(&mut self, word_size: usize, next: Option<&FreeLink>) {
        self.packed = 0;
        self.set_word_size(word_size);
        self.set_next(next);
    }

    fn get_next_offset(&self) -> u32 {
        self.packed & 0x7_ffff
    }

    fn get_next(&self) -> Option<&'static FreeLink> {
        let offset = self.get_next_offset() as isize;
        if offset == 0 { None } else { Some(self.jump(offset)) }
    }

    fn get_next_mut(&mut self) -> Option<&'static mut FreeLink> {
        let offset = self.get_next_offset() as isize;
        if offset == 0 { None } else { Some(self.jump_mut(offset)) }
    }

    fn set_next_offset(&mut self, offset: u32) {
        self.packed = (self.packed & 0xfff8_0000) | (offset & 0x7_ffff);
    }

    fn set_next(&mut self, next: Option<&FreeLink>) {
        self.set_next_offset(next.map(|next| self.compute_jump_to(next)).unwrap_or(0) as u32);
    }

    fn get_word_size(&self) -> usize {
        ((self.packed & 0xfff8_0000) >> 19) as usize
    }

    fn set_word_size(&mut self, word_size: usize) {
        self.packed = (self.packed & 0x7_ffff) | ((word_size << 19) as u32);
    }

    fn next_is_less_than(&self, link: &mut FreeLink) -> bool {
        self.get_next().map(|next| (next as *const FreeLink) < (link as *const FreeLink)).unwrap_or(false)
    }

    fn end_ptr(&self) -> *const FreeLink {
        (self as *const FreeLink).wrapping_add(self.get_word_size())
    }

    // chop off anything past `word_size`
    fn split_remainder(&mut self, word_size: usize) -> Option<&'static FreeLink> {
        let remain = self.get_word_size() - word_size;
        if remain == 0 { return None }
        let next = self.jump_mut(word_size as isize);
        next.init(remain, self.get_next());
        Some(next)
    }

    // merge this free block with the next, if they're right next to each other.
    fn check_merge_next(&mut self) {
        if let Some(next) = self.get_next() && ptr::eq(next, self.end_ptr()) {
            self.set_word_size(self.get_word_size() + next.get_word_size());
            if next.get_next_offset() == 0 {
                self.set_next_offset(0);
            } else {
                self.set_next_offset(self.get_next_offset() + next.get_next_offset());
            }
        }
    }
}


// hide the mess of dealing with the free list, which is kept sorted by
// pointer to allow compaction.
#[derive(Debug)]
pub struct FreeList {
    start: FreeLink,
}

impl FreeList {
    fn init(&mut self) {
        self.start.packed = 0;
    }

    fn total_word_count(&self) -> (usize, usize) {
        let mut ret = 0;
        let mut count = 0;
        let mut cur = &self.start;
        while let Some(next) = cur.get_next() {
            count += 1;
            ret += cur.get_word_size();
            cur = next;
        }
        ret += cur.get_word_size();
        (count, ret)
    }

    // keep them sorted in a single-linked list for compaction
    fn insert(&mut self, ptr: *mut u32, word_size: usize) {
        let link = FreeLink::bless(ptr);
        let mut cur = &mut self.start;
        while cur.next_is_less_than(link) { cur = cur.get_next_mut().unwrap() }

        if ptr::eq(cur.end_ptr(), link) {
            // merge with previous
            cur.set_word_size(cur.get_word_size() + word_size);
            cur.check_merge_next();
        } else {
            link.init(word_size, cur.get_next());
            cur.set_next(Some(link));
            link.check_merge_next();
        }
    }

    fn remove_last(&mut self) -> Option<&FreeLink> {
        let mut cur = &mut self.start;
        while let Some(next) = cur.get_next_mut() {
            if next.get_next().is_none() { break }
            cur = next;
        }
        let last = cur.get_next();
        cur.set_next(None);
        last
    }

    fn find(&mut self, word_size: usize, align_bits: usize) -> Option<*mut u32> {
        let mut cur = Some(&mut self.start);
        let mut prev: Option<&mut FreeLink> = None;
        while let Some(link) = cur.take() {
            if link.get_word_size() >= word_size && align_of(link as *mut FreeLink) >= align_bits {
                cur = Some(link);
                break;
            }
            let next = link.get_next_mut();
            prev.replace(link);
            cur = next;
        }
        let link = cur?;

        // next block is either remainder of this one, or the original next-block
        let next_block = link.split_remainder(word_size).or_else(|| link.get_next());
        // `prev` is always set here, because the `self.start` link is always
        // 0-length, requiring at least one trip around the loop.
        prev.unwrap().set_next(next_block);
        Some(link as *mut FreeLink as *mut u32)
    }
}

impl fmt::Display for FreeList {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "FreeList(")?;
        let mut cur = &self.start;
        while cur.get_next_offset() > 0 {
            write!(f, "${:x} -> @${:x}=", cur.get_word_size(), cur.get_next_offset())?;
            cur = cur.get_next().unwrap();
        }
        write!(f, "${:x})", cur.get_word_size())
    }
}


/// Simple memory allocator. It manages a region of memory defined by either
/// an array or a pair of pointers (begin, end), and hands out blocks in
/// (wasm-sized, 32 bit) words, using a bump pointer and a free list. This is
/// a bog-standard algorithm I remember from school in the 1990s, so nothing
/// fancy.
///
/// - New allocations will use the first block on the free list that has
///   space and fits the requested alignment. If none of those work, the
///   bump pointer is used until it hits the end of the region, after which
///   allocations fail.
/// - Freed blocks are inserted into the free list, maintaining address order.
///   Neighbor blocks are merged. Any freed block that touches the bump
///   pointer is given back to the bump pointer (and removed from the free
///   list).
///
/// This allocator may also be used as part of a larger allocation strategy.
/// For example, separate SimpleAlloc objects could manage separate regions
/// based on block size.
///
/// Each SimpleAllocator can manage up to 2MB, and allocate blocks up to
/// 32KB each. This is a limitation of the compact storage of the free list.
/// To manage more memory, chain multiple SimpleAllocator together. Large
/// allocations (>= 32KB) should be handled separately.
///
/// NOT THREAD-SAFE.
#[derive(Debug)]
pub struct SimpleAllocator {
    pub region_next: *mut SimpleAllocator,   // for use by a meta-allocator
    pub region_end: *mut u32,
    pub unused_start: *mut u32,
    pub free_list: FreeList,
}

impl fmt::Display for SimpleAllocator {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f, "SimpleAllocator(@${:08x}:${:08x}, unused=@${:08x}, {}, free_list={})",
            self.region_start() as usize, self.region_end as usize,
            self.unused_start as usize, self.get_stats(), self.free_list
        )?;
        if let Some(next) = self.get_next_region() {
            write!(f, " -> {}", next)?;
        }
        Ok(())
    }
}

impl SimpleAllocator {
    // use an array slice as allocatable memory
    pub fn init_at_mem<T>(mem: &mut [T]) -> &mut SimpleAllocator {
        let region_start = mem as *mut [T] as *mut u32;
        let region_end = (&mut mem[0] as *mut T).wrapping_add(mem.len()) as *mut u32;
        unsafe { SimpleAllocator::init_at(region_start, region_end) }
    }

    /// # Safety
    /// This is intrinsically unsafe. It's taking a block of memory and
    /// turning it into an allocation pool for use as rust's system allocator
    /// under wasm.
    pub unsafe fn init_at(region_start: *mut u32, region_end: *mut u32) -> &'static mut SimpleAllocator {
        let alloc = unsafe { &mut *(region_start as *mut SimpleAllocator) };
        let region_start = unsafe { (alloc as *mut SimpleAllocator).offset(1) as *mut u32 };
        alloc.init(region_start, region_end);
        alloc
    }

    fn init(&mut self, region_start: *mut u32, region_end: *mut u32) {
        assert!((region_end as usize) - (region_start as usize) <= MAX_HEAP);
        self.region_next = ptr::null_mut();
        self.region_end = region_end;
        self.unused_start = region_start;
        self.free_list.init();
    }

    pub fn get_next_region(&self) -> Option<&SimpleAllocator> {
        if self.region_next.is_null() { return None }
        Some(unsafe { &mut *self.region_next })
    }

    pub fn get_next_region_mut(&mut self) -> Option<&mut SimpleAllocator> {
        if self.region_next.is_null() { return None }
        Some(unsafe { &mut *self.region_next })
    }

    pub fn set_next_region(&mut self, next: Option<&mut SimpleAllocator>) {
        self.region_next = next.map(|next| next as *mut SimpleAllocator).unwrap_or(ptr::null_mut());
    }

    /// # Safety
    /// This is intrinsically unsafe. It extends the region of memory managed
    /// by this allocator. Use this when the heap can grow (via unix `brk()`
    /// or wasm `memory_grow()`).
    pub unsafe fn extend(&mut self, region_end: *mut u32) {
        assert!((region_end as usize) - (self.region_start() as usize) <= MAX_HEAP);
        self.region_end = region_end;
    }

    pub fn region_start(&self) -> *const u32 {
        (self as *const SimpleAllocator).wrapping_add(1) as *const u32
    }

    pub fn region_size(&self) -> usize {
        unsafe { (self.region_end as *mut u8).offset_from(self.region_start() as *const u8) as usize }
    }

    pub fn get_stats(&self) -> SimpleAllocatorStats {
        let total_words = unsafe { self.region_end.offset_from(self.region_start()) } as usize;
        let allocated_words = unsafe { self.unused_start.offset_from(self.region_start()) } as usize;
        let unused_words = total_words - allocated_words;
        let (fragments, fragment_words) = self.free_list.total_word_count();
        let mut stats = SimpleAllocatorStats {
            regions: 1,
            total_words,
            allocated_words,
            unused_words,
            fragment_words,
            fragments,
        };
        if let Some(next) = self.get_next_region() {
            stats.add(next.get_stats());
        }
        stats
    }

    pub fn dealloc(&mut self, ptr: *mut u32, word_size: usize) {
        if (ptr as *const _) < self.region_start() || ptr >= self.region_end {
            // not in our region? try the next in the chain, if it exists
            if let Some(next) = self.get_next_region_mut() {
                next.dealloc(ptr, word_size);
            }
            return;
        }

        self.free_list.insert(ptr, word_size);
        if ptr.wrapping_add(word_size) == self.unused_start {
            // pull the last block (which may have been merged with previous) into the unused region
            if let Some(last) = self.free_list.remove_last() {
                self.unused_start = last as *const FreeLink as *mut u32;
            }
        }
    }

    // align_bits: # of trailing zero bits in the address
    fn alloc_in_region(&mut self, word_size: usize, align_bits: usize) -> Option<*mut u32> {
        self.free_list.find(word_size, align_bits).or_else(|| {
            // alloc new
            let align_mask = (1 << align_bits) - 1;
            let start = self.unused_start as usize;
            let start_align = (start + align_mask) & !align_mask;

            // can save a lot of effort if there isn't actually enough room left
            if (((self.region_end as usize) - start_align) >> 2) < word_size { return None }

            if start != start_align {
                // skip words for alignment: can only happen for align_bits > 2
                let word_size = (start_align - start) >> 2;
                assert!(word_size > 0);
                self.free_list.insert(self.unused_start, word_size);
                self.unused_start = self.unused_start.wrapping_add(word_size);
            }

            let ret = Some(self.unused_start);
            self.unused_start = self.unused_start.wrapping_add(word_size);
            ret
        })
    }

    pub fn alloc(&mut self, word_size: usize, align_bits: usize) -> Option<*mut u32> {
        if word_size > (MAX_ALLOC >> 2) { return None }
        self.alloc_in_region(word_size, align_bits).or_else(|| {
            // try the next region in the chain, if it exists
            if let Some(next) = self.get_next_region_mut() {
                next.alloc(word_size, align_bits)
            } else {
                None
            }
        })
    }
}


#[inline(always)]
fn align_of<T>(ptr: *mut T) -> usize {
    (ptr as usize).trailing_zeros() as usize
}


#[cfg(test)]
mod tests {
    use core::mem::size_of;
    use super::*;


    #[test]
    fn setup() {
        let mut mem: [u32; 128] = [0; 128];
        let mem_addr = &mem as *const u32 as usize;
        let alloc = SimpleAllocator::init_at_mem(&mut mem);
        assert_eq!(alloc.region_next, ptr::null_mut());
        assert_eq!(alloc.region_end as usize, mem_addr + 128 * size_of::<u32>());
        assert_eq!(alloc.unused_start as usize, mem_addr + size_of::<SimpleAllocator>());
        assert_eq!(alloc.free_list.start.packed, 0);
    }

    #[test]
    fn alloc_dealloc_restores_state() {
        let mut mem: [u32; 128] = [0; 128];
        let alloc = SimpleAllocator::init_at_mem(&mut mem);
        let start = alloc.unused_start;

        assert_eq!(alloc.alloc(3, 0), Some(start));
        assert_eq!(alloc.unused_start, start.wrapping_add(3));
        assert_eq!(alloc.free_list.start.packed, 0);

        alloc.dealloc(start, 3);
        assert_eq!(alloc.unused_start, start);
        assert_eq!(alloc.free_list.start.packed, 0);
    }

    #[test]
    fn alloc_dealloc_realloc() {
        let mut mem: [u32; 128] = [0; 128];
        let alloc = SimpleAllocator::init_at_mem(&mut mem);
        let start = alloc.unused_start;

        // allocate two blocks
        assert_eq!(alloc.alloc(3, 0), Some(start));
        assert_eq!(alloc.unused_start, start.wrapping_add(3));
        assert_eq!(alloc.free_list.start.packed, 0);

        assert_eq!(alloc.alloc(3, 0), Some(start.wrapping_add(3)));
        assert_eq!(alloc.unused_start, start.wrapping_add(6));
        assert_eq!(alloc.free_list.start.packed, 0);

        // free the first block
        alloc.dealloc(start, 3);
        assert_eq!(alloc.unused_start, start.wrapping_add(6));
        assert_eq!(alloc.free_list.start.packed, 2);

        // grab the first block again
        assert_eq!(alloc.alloc(3, 0), Some(start));
        assert_eq!(alloc.unused_start, start.wrapping_add(6));
        assert_eq!(alloc.free_list.start.packed, 0);
    }

    #[test]
    fn coalesce_dealloc() {
        let mut mem: [u32; 128] = [0; 128];
        let alloc = SimpleAllocator::init_at_mem(&mut mem);
        let start = alloc.unused_start;

        let block1 = alloc.alloc(3, 0).unwrap();
        let block2 = alloc.alloc(3, 0).unwrap();
        let block3 = alloc.alloc(3, 0).unwrap();
        assert_eq!(block1, start);
        assert_eq!(block2, start.wrapping_add(3));
        assert_eq!(block3, start.wrapping_add(6));
        assert_eq!(alloc.unused_start, start.wrapping_add(9));
        assert_eq!(alloc.free_list.start.packed, 0);

        alloc.dealloc(block1, 3);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 9);
        assert_eq!(stats.fragment_words, 3);
        assert_eq!(stats.fragments, 1);
        alloc.dealloc(block2, 3);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 9);
        assert_eq!(stats.fragment_words, 6);
        assert_eq!(stats.fragments, 1);
        alloc.dealloc(block3, 3);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 0);
        assert_eq!(stats.fragment_words, 0);
        assert_eq!(stats.fragments, 0);
    }

    #[test]
    fn coalesce_middle() {
        let mut mem: [u32; 128] = [0; 128];
        let alloc = SimpleAllocator::init_at_mem(&mut mem);
        let start = alloc.unused_start;

        let block1 = alloc.alloc(3, 0).unwrap();
        let block2 = alloc.alloc(3, 0).unwrap();
        let block3 = alloc.alloc(3, 0).unwrap();
        assert_eq!(block1, start);
        assert_eq!(block2, start.wrapping_add(3));
        assert_eq!(block3, start.wrapping_add(6));
        assert_eq!(alloc.unused_start, start.wrapping_add(9));
        assert_eq!(alloc.free_list.start.packed, 0);

        alloc.dealloc(block1, 3);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 9);
        assert_eq!(stats.fragment_words, 3);
        assert_eq!(stats.fragments, 1);
        alloc.dealloc(block3, 3);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 6);
        assert_eq!(stats.fragment_words, 3);
        assert_eq!(stats.fragments, 1);
        alloc.dealloc(block2, 3);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 0);
        assert_eq!(stats.fragment_words, 0);
        assert_eq!(stats.fragments, 0);
    }

    #[test]
    fn merge_forward() {
        let mut mem: [u32; 128] = [0; 128];
        let alloc = SimpleAllocator::init_at_mem(&mut mem);

        let block1 = alloc.alloc(3, 0).unwrap();
        let block2 = alloc.alloc(3, 0).unwrap();
        let block3 = alloc.alloc(3, 0).unwrap();

        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 9);
        assert_eq!(stats.fragment_words, 0);
        assert_eq!(stats.fragments, 0);
        alloc.dealloc(block2, 3);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 9);
        assert_eq!(stats.fragment_words, 3);
        assert_eq!(stats.fragments, 1);
        alloc.dealloc(block1, 3);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 9);
        assert_eq!(stats.fragment_words, 6);
        assert_eq!(stats.fragments, 1);
        alloc.dealloc(block3, 3);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 0);
        assert_eq!(stats.fragment_words, 0);
        assert_eq!(stats.fragments, 0);
    }

    #[test]
    fn split_remainder() {
        let mut mem: [u32; 128] = [0; 128];
        let alloc = SimpleAllocator::init_at_mem(&mut mem);

        let block1 = alloc.alloc(4, 0).unwrap();
        let block2 = alloc.alloc(4, 0).unwrap();
        alloc.dealloc(block1, 4);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 8);
        assert_eq!(stats.fragment_words, 4);
        assert_eq!(stats.fragments, 1);

        let block3 = alloc.alloc(2, 0).unwrap();
        assert_eq!(block1, block3);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 8);
        assert_eq!(stats.fragment_words, 2);
        assert_eq!(stats.fragments, 1);

        alloc.dealloc(block2, 4);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 2);
        assert_eq!(stats.fragment_words, 0);
        assert_eq!(stats.fragments, 0);
    }

    #[test]
    fn alignment() {
        let mut mem: [u32; 128] = [0; 128];
        let alloc = SimpleAllocator::init_at_mem(&mut mem);
        let start = alloc.unused_start;

        let mut block = alloc.alloc(1, 0).unwrap();
        let mut count = 1;
        while (block as usize) & 0x1f != 0 {
            block = alloc.alloc(1, 0).unwrap();
            count += 1;
        }

        // insist on 32-byte alignment, meaning it has to skip (fragment) 28
        let padding = alloc.unused_start;
        let block = alloc.alloc(1, 5).unwrap();
        assert_eq!(block as usize, (padding as usize) + 28);
        let stats = alloc.get_stats();
        assert_eq!(stats.fragment_words, 7);
        assert_eq!(stats.fragments, 1);

        alloc.dealloc(start, count);
        alloc.dealloc(block, 1);
        let stats = alloc.get_stats();
        assert_eq!(stats.allocated_words, 0);
        assert_eq!(stats.fragment_words, 0);
        assert_eq!(stats.fragments, 0);
    }
}