142 lines
3.9 KiB
Rust
142 lines
3.9 KiB
Rust
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);
|
|
}
|
|
}
|