Newer
Older
vengine / engine-core / src / global_allocator.rs
Maple Nebel Maple Nebel 9 days ago 2 KB things
#[cfg(debug_assertions)]
use std::sync::atomic::AtomicUsize;
use std::{
    alloc::{GlobalAlloc, System},
    ptr::null_mut,
    sync::atomic::{AtomicPtr, Ordering::Relaxed},
};

struct AllocatorHandler {
    prim_entry: AtomicPtr<FreeBlockHeader>,
}

struct FreeBlockHeader {
    size: usize,
    next: AtomicPtr<FreeBlockHeader>,
}

const INITIAL_POOL_SIZE: usize = 0x10;
struct InitialFreePool {
    header: FreeBlockHeader,
    rest: [u8; INITIAL_POOL_SIZE],
}

static INITIAL_POOL: InitialFreePool = InitialFreePool {
    header: FreeBlockHeader {
        size: INITIAL_POOL_SIZE,
        next: AtomicPtr::new(null_mut()),
    },
    rest: [0; _],
};

/*
thread_local! {
    static ALLOCATOR_STACK: SmallVec<[15;]>
}*/

#[global_allocator]
static GLOBAL_ALLOCATOR_HANDLER: AllocatorHandler = AllocatorHandler {
    prim_entry: AtomicPtr::new((&INITIAL_POOL.header) as *const _ as *mut _),
};

#[cfg(debug_assertions)]
pub static ALLOCATION_SIZE_ORDER: [AtomicUsize; 64] = unsafe { std::mem::zeroed() };

unsafe impl GlobalAlloc for AllocatorHandler {
    unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
        if let Some(a) = ALLOCATION_SIZE_ORDER.get(layout.size().ilog2() as usize) {
            a.fetch_add(1, Relaxed);
        }
        System.alloc(layout)
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
        if let Some(a) = ALLOCATION_SIZE_ORDER.get(layout.size().ilog2() as usize) {
            a.fetch_sub(1, Relaxed);
        }
        System.dealloc(ptr, layout);
    }
    unsafe fn alloc_zeroed(&self, layout: std::alloc::Layout) -> *mut u8 {
        if let Some(a) = ALLOCATION_SIZE_ORDER.get(layout.size().ilog2() as usize) {
            a.fetch_add(1, Relaxed);
        }
        System.alloc_zeroed(layout)
    }
    unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
        if let Some(a) = ALLOCATION_SIZE_ORDER.get(layout.size().ilog2() as usize) {
            a.fetch_sub(1, Relaxed);
        }
        if let Some(a) = ALLOCATION_SIZE_ORDER.get(new_size.ilog2() as usize) {
            a.fetch_add(1, Relaxed);
        }
        System.realloc(ptr, layout, new_size)
    }
}