alloc_cat 1.1.1

use core::alloc::{GlobalAlloc, Layout};
use core::cell::{Cell, RefCell};
use core::fmt;
use core::mem;
use core::ptr;

use super::{Heap, PAGE_SIZE};
use super::allocator_stats::{AllocatorStats, Statsable};
use super::metrics::Metrics;
use super::simple_allocator::{MAX_ALLOC, MAX_HEAP, SimpleAllocator};


// linked list of large (>32KB) allocations, which are just sets of full
// 64KB pages. the end of each allocation is a PageLink to the next one,
// until `pages` is null.
#[derive(Copy, Clone, Debug)]
pub struct PageLink {
    page_count: u16,
    unused_words: u16,  // words on the last page that weren't requested
    pages: *mut u8,
}

// `unused_words` must always be at least size_of<PageLink> smaller than a full page
const PAGES_ARE_FREE: u16 = u16::MAX;

impl PageLink {
    pub const fn blank() -> PageLink {
        PageLink { page_count: 0, unused_words: PAGES_ARE_FREE, pages: ptr::null_mut() }
    }

    // turn a pointer-to-pages into a &mut to the PageLink at the end of them
    fn link_for(block: *mut u8, page_count: u16) -> &'static mut PageLink {
        let link = (block.wrapping_add((page_count as usize) * PAGE_SIZE) as *mut PageLink).wrapping_sub(1);
        unsafe { &mut *link }
    }

    fn next(&self) -> Option<&'static PageLink> {
        if self.pages.is_null() { return None }
        Some(PageLink::link_for(self.pages, self.page_count))
    }

    fn next_mut(&mut self) -> Option<&'static mut PageLink> {
        if self.pages.is_null() { return None }
        Some(PageLink::link_for(self.pages, self.page_count))
    }

    fn in_use(&self) -> bool {
        self.unused_words != PAGES_ARE_FREE
    }

    fn set_free(&mut self) {
        self.unused_words = PAGES_ARE_FREE;
    }

    fn end_ptr(&self) -> *const u8 {
        (self as *const PageLink).wrapping_add(1) as *const u8
    }

    // merge this free block with the next, if they're right next to each other.
    fn check_merge_next(&mut self) {
        if self.in_use() { return; }
        if let Some(next) = self.next_mut() && !next.in_use() && !ptr::eq(self.end_ptr(), next.pages) {
            self.page_count += next.page_count;
        }
    }

    fn create_free(&mut self, block: *mut u8, page_count: u16) {
        self.page_count = page_count;
        self.unused_words = PAGES_ARE_FREE;
        self.pages = block;
    }

    fn max_words_available(&self) -> usize {
        ((self.page_count as usize) * PAGE_SIZE - mem::size_of::<PageLink>()) >> 2
    }

    fn words_in_use(&self) -> usize {
        if self.in_use() {
            self.max_words_available() - (self.unused_words as usize)
        } else {
            0
        }
    }

    fn set_used_for(&mut self, word_count: usize) {
        let unused_words = self.max_words_available() - word_count;
        assert!(unused_words < (PAGE_SIZE >> 2));
        self.unused_words = unused_words as u16;
    }
}

const fn page_count_for_words(word_count: usize) -> u16 {
    (((word_count << 2) + (PAGE_SIZE - 1) + mem::size_of::<PageLink>()) / PAGE_SIZE) as u16
}


/// A replacement global allocator for wasm. It uses a linked list of 64KB
/// pages for large allocations, and a linked list of SimpleAllocators for
/// smaller ones. (Each SimpleAllocator can manage up to 2MB.)
///
/// NOT THREAD-SAFE.
#[derive(Debug)]
pub struct Allocator<'a> {
    heap: &'a dyn Heap,
    alloc: Cell<*mut SimpleAllocator>,
    // large allocations:
    large_pages: RefCell<PageLink>,

    // tracking metrics:
    metrics: Metrics,
}

// wasm hack
unsafe impl Sync for Allocator<'_> {}

impl<'a> Allocator<'a> {
    pub const fn new(heap: &'a dyn Heap) -> Allocator<'a> {
        Allocator {
            heap,
            alloc: Cell::new(ptr::null_mut()),
            large_pages: RefCell::new(PageLink::blank()),
            metrics: Metrics::new(),
        }
    }

    fn allocator(&self) -> &'static mut SimpleAllocator {
        let alloc = self.alloc.get();
        if !alloc.is_null() {
            return unsafe { &mut *alloc };
        }

        // initial configuration: detect starting heap size and use that
        let region_start = (self.heap.heap_start() + 3) & !3;
        let mut region_end = self.heap.heap_end() & !3;
        assert!(region_end >= region_start);
        if region_end - region_start < 1024 {
            // the starting heap is too small to bother with: add a page.
            self.heap.heap_grow(1);
            region_end = self.heap.heap_end();
        }

        let alloc = unsafe { SimpleAllocator::init_at(region_start as *mut u32, region_end as *mut u32) };
        self.alloc.set(alloc as *mut _);
        alloc
    }

    fn create_region(&self) -> Option<&'static mut SimpleAllocator> {
        self.heap.heap_grow(1).map(|ptr| {
            let region_start = ptr as *mut u32;
            let region_end = self.heap.heap_end() as *mut u32;
            let alloc = unsafe { SimpleAllocator::init_at(region_start, region_end) };
            let next = self.alloc.get();
            let next = if next.is_null() { None } else { Some(unsafe { &mut *next }) };
            alloc.set_next_region(next);
            self.alloc.set(alloc as *mut _);
            alloc
        })
    }

    // use region allocator
    fn alloc_small(&self, word_count: usize, align_bits: usize) -> *mut u8 {
        let mut allocator = self.allocator();
        assert!(allocator.region_end as *const u32 >= allocator.region_start());

        loop {
            if let Some(block) = allocator.alloc(word_count, align_bits) {
                return block as *mut u8;
            }

            if allocator.region_size() + PAGE_SIZE <= MAX_HEAP &&
                (allocator.region_end as usize) == self.heap.heap_end() {
                // give it another page
                if self.heap.heap_grow(1).is_some() {
                    unsafe { allocator.extend(self.heap.heap_end() as *mut u32); }
                } else {
                    // FAIL
                    return ptr::null_mut();
                }
            } else {
                // create a new region
                if let Some(a) = self.create_region() {
                    allocator = a;
                } else {
                    // FAIL
                    return ptr::null_mut();
                }
            }
        }
    }

    fn alloc_large(&self, word_count: usize) -> *mut u8 {
        let page_count = page_count_for_words(word_count);
        let mut start = self.large_pages.borrow_mut();
        let mut next = Some(&mut *start);
        while let Some(link) = next {
            if !link.in_use() && link.page_count >= page_count {
                // score!
                if link.page_count > page_count {
                    // split off the last page(s)
                    let remain_page_count = link.page_count - page_count;
                    link.page_count = page_count;
                    let new_link = link.next_mut().unwrap();
                    new_link.create_free(link.pages.wrapping_add(PAGE_SIZE * (page_count as usize)), remain_page_count);
                }
                link.set_used_for(word_count);
                return link.pages;
            }
            next = link.next_mut();
        }

        // time to make the donuts.
        if let Some(pages) = self.heap.heap_grow(page_count as usize) {
            let link = PageLink::link_for(pages, page_count);
            link.clone_from(&start);
            start.create_free(pages, page_count);
            start.set_used_for(word_count);
            pages
        } else {
            // FAIL
            ptr::null_mut()
        }
    }

    fn dealloc_large(&self, ptr: *mut u8) {
        let mut start = self.large_pages.borrow_mut();
        let mut prev: Option<&mut PageLink> = None;
        let mut next = Some(&mut *start);
        while let Some(link) = next {
            if link.in_use() && link.pages == ptr {
                link.set_free();
                link.check_merge_next();
                if let Some(prev) = prev {
                    prev.check_merge_next();
                }
                return;
            }
            next = link.next_mut();
            prev.replace(link);
        }
    }
}

impl fmt::Display for Allocator<'_> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Allocator({}, {})", self.heap, self.allocator())
    }
}

impl Statsable for Allocator<'_> {
    fn get_stats(&self) -> AllocatorStats {
        let mut stats = AllocatorStats::from_simple_allocator_stats(self.allocator().get_stats());
        self.metrics.get_stats(&mut stats);

        let start = self.large_pages.borrow();
        let mut next = Some(&*start);
        while let Some(link) = next {
            stats.total_large_pages += link.page_count as usize;
            if !link.in_use() {
                stats.unused_large_pages += link.page_count as usize;
            } else {
                stats.allocated_large_pages_words += link.words_in_use();
            }
            next = link.next();
        }

        stats.heap_start = self.heap.heap_start();
        stats
    }

    fn reset_trip(&self) {
        self.metrics.reset_trip();
    }
}

unsafe impl GlobalAlloc for Allocator<'_> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let word_count = (layout.size() + 3) >> 2;
        let align_bits = layout.align().trailing_zeros() as usize;
        // ensure the region allocator exists:
        self.allocator();

        self.metrics.add(layout);

        if word_count > (MAX_ALLOC >> 2) {
            self.alloc_large(word_count)
        } else {
            self.alloc_small(word_count, align_bits)
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        self.metrics.subtract(layout);
        let word_count = (layout.size() + 3) >> 2;
        if word_count > (MAX_ALLOC >> 2) {
            self.dealloc_large(ptr);
        } else {
            self.allocator().dealloc(ptr as *mut u32, word_count);
        }
    }
}


#[cfg(test)]
mod tests {
    use ptr::slice_from_raw_parts_mut;

    use crate::MmapHeap;

    use super::*;


    static ALLOC_32: fn () -> Layout = || Layout::from_size_align(32, 1).unwrap();
    static ALLOC_1024: fn () -> Layout = || Layout::from_size_align(1024, 1).unwrap();
    static ALLOC_48K: fn () -> Layout = || Layout::from_size_align(48 * 1024, 1).unwrap();
    static ALLOC_96K: fn () -> Layout = || Layout::from_size_align(96 * 1024, 1).unwrap();


    #[test]
    fn simple() {
        let heap = MmapHeap::new(256 * 1024);
        let alloc = Allocator::new(&heap);
        let ptr = unsafe { alloc.alloc(ALLOC_32()) };
        let stats = alloc.get_stats();
        assert_eq!(stats.regions, 1);
        assert_eq!(stats.allocated_words, 8);
        assert_eq!(stats.allocated_large_pages_words, 0);
        unsafe { alloc.dealloc(ptr, ALLOC_32()); }
        let stats = alloc.get_stats();
        assert_eq!(stats.regions, 1);
        assert_eq!(stats.allocated_words, 0);
        assert_eq!(stats.allocated_large_pages_words, 0);
    }

    #[test]
    fn multiple_regions() {
        let heap = MmapHeap::new(256 * 1024);
        let alloc = Allocator::new(&heap);
        let _ptr = unsafe { alloc.alloc(ALLOC_1024()) };
        let stats = alloc.get_stats();
        assert_eq!(stats.regions, 1);
        assert_eq!(stats.allocated_words, 256);
        assert_eq!(stats.allocated_large_pages_words, 0);

        // alloc random page, messing up our precious contiguous region
        assert!(heap.heap_grow(1).is_some());
        for _ in 0..64 { unsafe { alloc.alloc(ALLOC_1024()); } }
        let stats = alloc.get_stats();
        assert_eq!(stats.regions, 2);
        assert_eq!(stats.allocated_words, 256 * 65);
        assert_eq!(stats.allocated_large_pages_words, 0);
    }

    #[test]
    fn large_simple() {
        let heap = MmapHeap::new(256 * 1024);
        let alloc = Allocator::new(&heap);
        let ptr = unsafe { alloc.alloc(ALLOC_96K()) };
        let stats = alloc.get_stats();
        assert_eq!(stats.regions, 1);
        assert_eq!(stats.allocated_words, 0);
        assert_eq!(stats.allocated_large_pages_words, 96 * 256);
        assert_eq!(stats.total_large_pages, 2);
        assert_eq!(stats.unused_large_pages, 0);

        unsafe { alloc.dealloc(ptr, ALLOC_96K()); }
        let stats = alloc.get_stats();
        assert_eq!(stats.regions, 1);
        assert_eq!(stats.allocated_words, 0);
        assert_eq!(stats.allocated_large_pages_words, 0);
        assert_eq!(stats.total_large_pages, 2);
        assert_eq!(stats.unused_large_pages, 2);
    }

    #[test]
    fn large_reuse() {
        let heap = MmapHeap::new(256 * 1024);
        let alloc = Allocator::new(&heap);
        let ptr1 = unsafe { alloc.alloc(ALLOC_96K()) };
        let _ptr2 = unsafe { alloc.alloc(ALLOC_48K()) };
        let stats = alloc.get_stats();
        assert_eq!(stats.regions, 1);
        assert_eq!(stats.allocated_words, 0);
        assert_eq!(stats.allocated_large_pages_words, 144 * 256);
        assert_eq!(stats.total_large_pages, 3);
        assert_eq!(stats.unused_large_pages, 0);

        unsafe { alloc.dealloc(ptr1, ALLOC_96K()); }
        let stats = alloc.get_stats();
        assert_eq!(stats.regions, 1);
        assert_eq!(stats.allocated_words, 0);
        assert_eq!(stats.allocated_large_pages_words, 48 * 256);
        assert_eq!(stats.total_large_pages, 3);
        assert_eq!(stats.unused_large_pages, 2);

        let _ptr1 = unsafe { alloc.alloc(ALLOC_96K()) };
        let stats = alloc.get_stats();
        assert_eq!(stats.regions, 1);
        assert_eq!(stats.allocated_words, 0);
        assert_eq!(stats.allocated_large_pages_words, 144 * 256);
        assert_eq!(stats.total_large_pages, 3);
        assert_eq!(stats.unused_large_pages, 0);
    }

    #[test]
    fn large_threshold() {
        let size1 = Layout::from_size_align(64 * 1024 - 4, 1).unwrap();
        let size2 = Layout::from_size_align(128 * 1024 - 4, 1).unwrap();

        let heap = MmapHeap::new(512 * 1024);
        let alloc = Allocator::new(&heap);
        let ptr1 = unsafe { alloc.alloc(size1) };
        let ptr2 = unsafe { alloc.alloc(size2) };
        let stats = alloc.get_stats();
        assert_eq!(stats.regions, 1);
        assert_eq!(stats.allocated_words, 0);
        assert_eq!(stats.total_large_pages, 5);
        assert_eq!(stats.allocated_large_pages_words, 192 * 256 - 2);
        assert_eq!(stats.unused_large_pages, 0);

        // if it failed to allocate an extra page for ptr1, we'll corrupt the ptr2 link here:
        let mem1 = unsafe { &mut *slice_from_raw_parts_mut(ptr1, 64 * 1024 - 4) };
        for b in mem1 { *b = 0xff; }

        let size3 = Layout::from_size_align(192 * 1024 - mem::size_of::<PageLink>(), 1).unwrap();
        unsafe { alloc.dealloc(ptr2, size2); }
        let stats = alloc.get_stats();
        assert_eq!(stats.total_large_pages, 5);
        assert_eq!(stats.allocated_large_pages_words, 64 * 256 - 1);
        assert_eq!(stats.unused_large_pages, 3);

        let ptr3 = unsafe { alloc.alloc(size3) };
        assert_eq!(ptr2, ptr3);
        let stats = alloc.get_stats();
        assert_eq!(stats.total_large_pages, 5);
        assert_eq!(stats.allocated_large_pages_words, 256 * 256 - 1 - (mem::size_of::<PageLink>() >> 2));
    }

    #[test]
    fn large_split() {
        let size_192k = Layout::from_size_align(192 * 1024 - 128, 1).unwrap();
        let size_64k = Layout::from_size_align(64 * 1024 - 128, 1).unwrap();
        let heap = MmapHeap::new(512 * 1024);
        let alloc = Allocator::new(&heap);

        // alloc 3 pages, then free them
        let ptr1 = unsafe { alloc.alloc(size_192k) };
        let stats = alloc.get_stats();
        assert_eq!(stats.total_large_pages, 3);
        assert_eq!(stats.allocated_large_pages_words, size_192k.size() >> 2);
        assert_eq!(stats.unused_large_pages, 0);
        unsafe { alloc.dealloc(ptr1, size_192k); }
        let stats = alloc.get_stats();
        assert_eq!(stats.total_large_pages, 3);
        assert_eq!(stats.allocated_large_pages_words, 0);
        assert_eq!(stats.unused_large_pages, 3);

        // alloc 1 page, 3 times
        let ptr1 = unsafe { alloc.alloc(size_64k) };
        let stats = alloc.get_stats();
        assert_eq!(stats.total_large_pages, 3);
        assert_eq!(stats.allocated_large_pages_words, size_64k.size() >> 2);
        assert_eq!(stats.unused_large_pages, 2);
        let ptr2 = unsafe { alloc.alloc(size_64k) };
        let stats = alloc.get_stats();
        assert_eq!(stats.total_large_pages, 3);
        assert_eq!(stats.allocated_large_pages_words, (size_64k.size() * 2) >> 2);
        assert_eq!(stats.unused_large_pages, 1);
        let ptr3 = unsafe { alloc.alloc(size_64k) };
        let stats = alloc.get_stats();
        assert_eq!(stats.total_large_pages, 3);
        assert_eq!(stats.allocated_large_pages_words, (size_64k.size() * 3) >> 2);
        assert_eq!(stats.unused_large_pages, 0);

        // free first and last page
        unsafe { alloc.dealloc(ptr1, size_64k); }
        let stats = alloc.get_stats();
        assert_eq!(stats.total_large_pages, 3);
        assert_eq!(stats.allocated_large_pages_words, (size_64k.size() * 2) >> 2);
        assert_eq!(stats.unused_large_pages, 1);
        unsafe { alloc.dealloc(ptr3, size_64k); }
        let stats = alloc.get_stats();
        assert_eq!(stats.total_large_pages, 3);
        assert_eq!(stats.allocated_large_pages_words, size_64k.size() >> 2);
        assert_eq!(stats.unused_large_pages, 2);

        // free the middle page. should be merged with the rest
        unsafe { alloc.dealloc(ptr2, size_64k); }
        let stats = alloc.get_stats();
        assert_eq!(stats.total_large_pages, 3);
        assert_eq!(stats.allocated_large_pages_words, 0);
        assert_eq!(stats.unused_large_pages, 3);

        // now grab the unified 3 pages again
        let _ptr1 = unsafe { alloc.alloc(size_192k) };
        let stats = alloc.get_stats();
        assert_eq!(stats.total_large_pages, 3);
        assert_eq!(stats.allocated_large_pages_words, size_192k.size() >> 2);
        assert_eq!(stats.unused_large_pages, 0);
    }

    #[test]
    fn large_then_small() {
        let size_15k = Layout::from_size_align(15 * 1024, 1).unwrap();
        let size_16k = Layout::from_size_align(16 * 1024, 1).unwrap();
        let size_64k = Layout::from_size_align(64 * 1024 - 128, 1).unwrap();
        let heap = MmapHeap::new(512 * 1024);
        let alloc = Allocator::new(&heap);

        assert_eq!(heap.heap_end() - heap.heap_start(), 0);
        let _ptr1 = unsafe { alloc.alloc(size_64k) };
        let stats = alloc.get_stats();
        assert_eq!(stats.total_large_pages, 1);
        assert_eq!(stats.allocated_large_pages_words, size_64k.size() >> 2);
        assert_eq!(stats.unused_large_pages, 0);
        assert_eq!(heap.heap_end() - heap.heap_start(), 128 * 1024);

        // now fill the single SimpleAllocator page
        for _i in 0..3 {
            unsafe { alloc.alloc(size_16k) };
        }
        unsafe { alloc.alloc(size_15k) };
        assert_eq!(heap.heap_end() - heap.heap_start(), 128 * 1024);

        // new 15k allocation will require a 2nd region
        let _ptr = unsafe { alloc.alloc(size_15k) };
        let stats = alloc.get_stats();
        assert_eq!(stats.regions, 2);
        assert_eq!(stats.allocated_words, 78 * 256);
        assert_eq!(stats.allocated_large_pages_words, size_64k.size() >> 2);
        assert_eq!(heap.heap_end() - heap.heap_start(), 192 * 1024);
    }

    #[test]
    fn trip_odometer() {
        let heap = MmapHeap::new(256 * 1024);
        let alloc = Allocator::new(&heap);
        let ptr = unsafe { alloc.alloc(ALLOC_1024()) };
        let stats = alloc.get_stats();
        assert_eq!(stats.high_water, 1024);
        assert_eq!(stats.trip_high_water, 1024);
        unsafe { alloc.dealloc(ptr, ALLOC_1024()) };

        alloc.reset_trip();
        let _ptr = unsafe { alloc.alloc(ALLOC_32()) };
        let stats = alloc.get_stats();
        assert_eq!(stats.high_water, 1024);
        assert_eq!(stats.trip_high_water, 32);
    }
}