mirmalloc 0.2.0

use std::alloc::{GlobalAlloc, Layout};
use std::cell::RefCell;
use std::sync::OnceLock;
use crate::heap::RawHeap;

pub struct ThreadPerCore;

// Thread-local heap, only initialized if init_heap_for_thread() is called
thread_local! {
    static THREAD_HEAP: RefCell<Option<RawHeap>> = RefCell::new(None);
}

static FALLBACK_HEAP: OnceLock<RawHeap> = OnceLock::new();

impl ThreadPerCore {
    /// Initialize a dedicated heap for the current thread.
    ///
    /// Threads that call this will use their own heap instead of the shared static heap.
    pub fn init_heap() {
        THREAD_HEAP.with(|heap| {
            *heap.borrow_mut() = Some(RawHeap::new());
        });
    }

    /// Returns the fallback heap, initializing it if necessary.
    #[inline]
    fn fallback() -> &'static RawHeap {
        FALLBACK_HEAP.get_or_init(RawHeap::new)
    }
}

unsafe impl GlobalAlloc for ThreadPerCore {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        match THREAD_HEAP.try_with(|c| c.borrow().as_ref().map(|h| unsafe { h.alloc(layout) })) {
            Ok(Some(ptr)) => ptr,
            _ => unsafe { Self::fallback().alloc(layout) },
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        match THREAD_HEAP.try_with(|c| c.borrow().as_ref().map(|h| unsafe { h.dealloc(ptr, layout) })) {
            Ok(Some(())) => {}
            _ => unsafe { Self::fallback().dealloc(ptr, layout) },
        }
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        match THREAD_HEAP.try_with(|c| c.borrow().as_ref().map(|h| unsafe { h.alloc_zeroed(layout) })) {
            Ok(Some(ptr)) => ptr,
            _ => unsafe { Self::fallback().alloc_zeroed(layout) },
        }
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        match THREAD_HEAP.try_with(|c| c.borrow().as_ref().map(|h| unsafe { h.realloc(ptr, layout, new_size) })) {
            Ok(Some(ptr)) => ptr,
            _ => unsafe { Self::fallback().realloc(ptr, layout, new_size) },
        }
    }
}