k-os/crates/libkos/src/alloc_impl.rs

use core::alloc::{GlobalAlloc, Layout};
use core::cell::UnsafeCell;
use core::ptr;

const PAGE_SIZE: usize = 4096;
const MAX_PAGES: usize = 64;

struct FreeBlock {
    size: usize,
    next: *mut FreeBlock,
}

pub struct KosAllocator {
    inner: UnsafeCell<AllocatorInner>,
}

struct AllocatorInner {
    free_list: *mut FreeBlock,
    pages: [usize; MAX_PAGES],
    page_count: usize,
}

unsafe impl Sync for KosAllocator {}

impl KosAllocator {
    pub const fn new() -> Self {
        Self {
            inner: UnsafeCell::new(AllocatorInner {
                free_list: ptr::null_mut(),
                pages: [0; MAX_PAGES],
                page_count: 0,
            }),
        }
    }

    fn inner_mut(&self) -> &mut AllocatorInner {
        unsafe { &mut *self.inner.get() }
    }
}

impl AllocatorInner {
    fn request_page(&mut self) -> *mut u8 {
        if self.page_count >= MAX_PAGES {
            return ptr::null_mut();
        }

        let count = match types::bounded::PageCount::new(1) {
            Some(c) => c,
            None => return ptr::null_mut(),
        };
        let va = match crate::syscall::request_pages(count) {
            Some(va) => va,
            None => return ptr::null_mut(),
        };

        self.pages[self.page_count] = va.raw();
        self.page_count += 1;
        va.raw() as *mut u8
    }

    fn alloc_from_free_list(&mut self, size: usize, align: usize) -> *mut u8 {
        let mut prev: *mut FreeBlock = ptr::null_mut();
        let mut current = self.free_list;

        while !current.is_null() {
            let block = unsafe { &mut *current };
            let addr = current as usize;
            let aligned = (addr + align - 1) & !(align - 1);
            let padding = aligned - addr;

            if block.size >= size + padding {
                let remaining = block.size - size - padding;

                if remaining >= core::mem::size_of::<FreeBlock>() {
                    let new_block = (aligned + size) as *mut FreeBlock;
                    unsafe {
                        (*new_block).size = remaining;
                        (*new_block).next = block.next;
                    }
                    if prev.is_null() {
                        self.free_list = new_block;
                    } else {
                        unsafe {
                            (*prev).next = new_block;
                        }
                    }
                } else {
                    if prev.is_null() {
                        self.free_list = block.next;
                    } else {
                        unsafe {
                            (*prev).next = block.next;
                        }
                    }
                }

                return aligned as *mut u8;
            }

            prev = current;
            current = block.next;
        }

        ptr::null_mut()
    }

    fn add_to_free_list(&mut self, ptr: *mut u8, size: usize) {
        let block = ptr as *mut FreeBlock;
        unsafe {
            (*block).size = size;
            (*block).next = self.free_list;
        }
        self.free_list = block;
    }
}

unsafe impl GlobalAlloc for KosAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let inner = self.inner_mut();
        let size = layout.size().max(core::mem::size_of::<FreeBlock>());
        let align = layout.align();

        let ptr = inner.alloc_from_free_list(size, align);
        if !ptr.is_null() {
            return ptr;
        }

        let pages_needed = (size + PAGE_SIZE - 1) / PAGE_SIZE;
        for _ in 0..pages_needed {
            let page = inner.request_page();
            if page.is_null() {
                return ptr::null_mut();
            }
            inner.add_to_free_list(page, PAGE_SIZE);
        }

        inner.alloc_from_free_list(size, align)
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let inner = self.inner_mut();
        let size = layout.size().max(core::mem::size_of::<FreeBlock>());
        inner.add_to_free_list(ptr, size);
    }
}