diff --git a/src/allocator.rs b/src/allocator.rs index 02bfc0a..91d6575 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -1,3 +1,6 @@ +pub mod fixed_size_block; +pub mod linked_list; +pub mod bump; use x86_64::{ structures::paging::{ mapper::MapToError, FrameAllocator, Mapper, Page, PageTableFlags, Size4KiB, @@ -6,13 +9,36 @@ use x86_64::{ }; use alloc::alloc::{GlobalAlloc, Layout}; use core::ptr::null_mut; -use linked_list_allocator::LockedHeap; +use fixed_size_block::FixedSizeBlockAllocator; + pub const HEAP_START: usize = 0x_4444_4444_0000; pub const HEAP_SIZE: usize = 100 * 1024; // 100 KiB #[global_allocator] -static ALLOCATOR: LockedHeap = LockedHeap::empty(); +static ALLOCATOR: Locked = Locked::new( + FixedSizeBlockAllocator::new()); +/// A wrapper around spin::Mutex to permit trait implementations. +pub struct Locked { + inner: spin::Mutex, +} + +impl Locked { + pub const fn new(inner: A) -> Self { + Locked { + inner: spin::Mutex::new(inner), + } + } + + pub fn lock(&self) -> spin::MutexGuard { + self.inner.lock() + } +} + +fn align_up(addr: usize, align: usize) -> usize { + (addr + align - 1) & !(align - 1) +} + pub struct Dummy; unsafe impl GlobalAlloc for Dummy { diff --git a/src/allocator/bump.rs b/src/allocator/bump.rs new file mode 100644 index 0000000..cf64201 --- /dev/null +++ b/src/allocator/bump.rs @@ -0,0 +1,56 @@ +use super::{align_up,Locked}; +use alloc::alloc::{GlobalAlloc, Layout}; +use core::ptr; + +pub struct BumpAllocator { + heap_start: usize, + heap_end: usize, + next: usize, + allocations: usize, +} + +impl BumpAllocator { + /// Creates a new empty bump allocator. + pub const fn new() -> Self { + BumpAllocator { + heap_start: 0, + heap_end: 0, + next: 0, + allocations: 0, + } + } + pub unsafe fn init(&mut self, heap_start: usize, heap_size: usize) { + self.heap_start = heap_start; + self.heap_end = heap_start + heap_size; + self.next = heap_start; + } +} + +unsafe impl GlobalAlloc for Locked { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let mut bump = self.lock(); // get a mutable reference + + let alloc_start = align_up(bump.next, layout.align()); + let alloc_end = match alloc_start.checked_add(layout.size()) { + Some(end) => end, + None => return ptr::null_mut(), + }; + + if alloc_end > bump.heap_end { + ptr::null_mut() // out of memory + } else { + bump.next = alloc_end; + bump.allocations += 1; + alloc_start as *mut u8 + } + } + + unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) { + let mut bump = self.lock(); // get a mutable reference + + bump.allocations -= 1; + if bump.allocations == 0 { + bump.next = bump.heap_start; + } + } +} \ No newline at end of file diff --git a/src/allocator/fixed_size_block.rs b/src/allocator/fixed_size_block.rs new file mode 100644 index 0000000..6d2a015 --- /dev/null +++ b/src/allocator/fixed_size_block.rs @@ -0,0 +1,90 @@ +use super::Locked; +use alloc::alloc::GlobalAlloc; +use alloc::alloc::Layout; +use core::{mem, ptr, ptr::NonNull}; +struct ListNode { + next: Option<&'static mut ListNode>, +} + +const BLOCK_SIZES: &[usize] = &[8, 16, 32, 64, 128, 256, 512, 1024, 2048]; + +pub struct FixedSizeBlockAllocator { + list_heads: [Option<&'static mut ListNode>; BLOCK_SIZES.len()], + fallback_allocator: linked_list_allocator::Heap, +} + +unsafe impl GlobalAlloc for Locked { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let mut allocator = self.lock(); + match FixedSizeBlockAllocator::list_index(&layout) { + Some(index) => { + match allocator.list_heads[index].take() { + Some(node) => { + allocator.list_heads[index] = node.next.take(); + node as *mut ListNode as *mut u8 + } + None => { + // no block exists in list => allocate new block + let block_size = BLOCK_SIZES[index]; + // only works if all block sizes are a power of 2 + let block_align = block_size; + let layout = Layout::from_size_align(block_size, block_align) + .unwrap(); + allocator.fallback_alloc(layout) + } + } + } + None => allocator.fallback_alloc(layout), + } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + let mut allocator = self.lock(); + match FixedSizeBlockAllocator::list_index(&layout) { + Some(index) => { + let new_node = ListNode { + next: allocator.list_heads[index].take(), + }; + // verify that block has size and alignment required for storing node + assert!(mem::size_of::() <= BLOCK_SIZES[index]); + assert!(mem::align_of::() <= BLOCK_SIZES[index]); + let new_node_ptr = ptr as *mut ListNode; + new_node_ptr.write(new_node); + allocator.list_heads[index] = Some(&mut *new_node_ptr); + } + None => { + let ptr = NonNull::new(ptr).unwrap(); + allocator.fallback_allocator.deallocate(ptr, layout); + } + } + } +} + +impl FixedSizeBlockAllocator { + /// Creates an empty FixedSizeBlockAllocator. + pub const fn new() -> Self { + const EMPTY: Option<&'static mut ListNode> = None; + FixedSizeBlockAllocator { + list_heads: [EMPTY; BLOCK_SIZES.len()], + fallback_allocator: linked_list_allocator::Heap::empty(), + } + } + + pub unsafe fn init(&mut self, heap_start: usize, heap_size: usize) { + self.fallback_allocator.init(heap_start, heap_size); + } + + /// Allocates using the fallback allocator. + fn fallback_alloc(&mut self, layout: Layout) -> *mut u8 { + match self.fallback_allocator.allocate_first_fit(layout) { + Ok(ptr) => ptr.as_ptr(), + Err(_) => ptr::null_mut(), + } + } + + /// Returns an index into the `BLOCK_SIZES` array. + fn list_index(layout: &Layout) -> Option { + let required_block_size = layout.size().max(layout.align()); + BLOCK_SIZES.iter().position(|&s| s >= required_block_size) + } +} diff --git a/src/allocator/linked_list.rs b/src/allocator/linked_list.rs new file mode 100644 index 0000000..d8f5b26 --- /dev/null +++ b/src/allocator/linked_list.rs @@ -0,0 +1,132 @@ +use super::align_up; +use core::mem; +use super::Locked; +use alloc::alloc::{GlobalAlloc, Layout}; +use core::ptr; + +pub struct LinkedListAllocator { + head: ListNode, +} + +unsafe impl GlobalAlloc for Locked { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + // perform layout adjustments + let (size, align) = LinkedListAllocator::size_align(layout); + let mut allocator = self.lock(); + + if let Some((region, alloc_start)) = allocator.find_region(size, align) { + let alloc_end = alloc_start.checked_add(size).expect("overflow"); + let excess_size = region.end_addr() - alloc_end; + if excess_size > 0 { + allocator.add_free_region(alloc_end, excess_size); + } + alloc_start as *mut u8 + } else { + ptr::null_mut() + } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + // perform layout adjustments + let (size, _) = LinkedListAllocator::size_align(layout); + + self.lock().add_free_region(ptr as usize, size) + } +} + + +impl LinkedListAllocator { + /// Creates an empty LinkedListAllocator. + pub const fn new() -> Self { + Self { + head: ListNode::new(0), + } + } + pub unsafe fn init(&mut self, heap_start: usize, heap_size: usize) { + self.add_free_region(heap_start, heap_size); + } + unsafe fn add_free_region(&mut self, addr: usize, size: usize) { + assert_eq!(align_up(addr, mem::align_of::()), addr); + assert!(size >= mem::size_of::()); + + // create a new list node and append it at the start of the list + let mut node = ListNode::new(size); + node.next = self.head.next.take(); + let node_ptr = addr as *mut ListNode; + node_ptr.write(node); + self.head.next = Some(&mut *node_ptr) + } + + fn find_region(&mut self, size: usize, align: usize) + -> Option<(&'static mut ListNode, usize)> + { + // reference to current list node, updated for each iteration + let mut current = &mut self.head; + // look for a large enough memory region in linked list + while let Some(ref mut region) = current.next { + if let Ok(alloc_start) = Self::alloc_from_region(®ion, size, align) { + // region suitable for allocation -> remove node from list + let next = region.next.take(); + let ret = Some((current.next.take().unwrap(), alloc_start)); + current.next = next; + return ret; + } else { + // region not suitable -> continue with next region + current = current.next.as_mut().unwrap(); + } + } + + // no suitable region found + None + } + + fn alloc_from_region(region: &ListNode, size: usize, align: usize) + -> Result + { + let alloc_start = align_up(region.start_addr(), align); + let alloc_end = alloc_start.checked_add(size).ok_or(())?; + + if alloc_end > region.end_addr() { + // region too small + return Err(()); + } + + let excess_size = region.end_addr() - alloc_end; + if excess_size > 0 && excess_size < mem::size_of::() { + // rest of region too small to hold a ListNode (required because the + // allocation splits the region in a used and a free part) + return Err(()); + } + + // region suitable for allocation + Ok(alloc_start) + } + + fn size_align(layout: Layout) -> (usize, usize) { + let layout = layout + .align_to(mem::align_of::()) + .expect("adjusting alignment failed") + .pad_to_align(); + let size = layout.size().max(mem::size_of::()); + (size, layout.align()) + } +} +struct ListNode { + size: usize, + next: Option<&'static mut ListNode>, +} + +impl ListNode { + const fn new(size: usize) -> Self { + ListNode { size, next: None } + } + + fn start_addr(&self) -> usize { + self as *const Self as usize + } + + fn end_addr(&self) -> usize { + self.start_addr() + self.size + } +} + diff --git a/src/lib.rs b/src/lib.rs index 8998104..d675b4f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,6 +5,7 @@ #![reexport_test_harness_main = "test_main"] #![feature(abi_x86_interrupt)] #![feature(alloc_error_handler)] +#![feature(const_mut_refs)] #[cfg(test)] use bootloader::{entry_point, BootInfo}; @@ -18,7 +19,7 @@ pub mod gdt; pub mod memory; pub mod allocator; -use core::{alloc::Layout, panic::PanicInfo}; +use core::panic::PanicInfo; pub trait Testable { fn run(&self) -> (); diff --git a/tests/heap_allocation.rs b/tests/heap_allocation.rs index 8e8f6c7..4b15401 100644 --- a/tests/heap_allocation.rs +++ b/tests/heap_allocation.rs @@ -61,3 +61,13 @@ fn many_boxes() { assert_eq!(*x, i); } } + +#[test_case] +fn many_boxes_long_lived() { + let long_lived = Box::new(1); // new + for i in 0..HEAP_SIZE { + let x = Box::new(i); + assert_eq!(*x, i); + } + assert_eq!(*long_lived, 1); // new +} \ No newline at end of file