feat: add fixed size block allocator

Author: criskell

kernel/src/allocator.rs

@@ -8,9 +8,13 @@ use x86_64::{
     },
 };
 
-use crate::allocator::{bump::BumpAllocator, linked_list::LinkedListAllocator};
+use crate::allocator::{
+    bump::BumpAllocator, fixed_size_block::FixedSizeBlockAllocator,
+    linked_list::LinkedListAllocator,
+};
 
 pub mod bump;
+pub mod fixed_size_block;
 pub mod linked_list;
 
 pub const HEAP_START: usize = 0x_4444_4444_0000;
@@ -49,7 +53,7 @@ pub fn init_heap(
 }
 
 #[global_allocator]
-static ALLOCATOR: LockedHeap = Locked::new(LinkedListAllocator::new());
+static ALLOCATOR: LockedHeap = Locked::new(FixedSizeBlockAllocator::new());
 
 pub struct Dummy;
 

kernel/src/allocator/fixed_size_block.rs

@@ -0,0 +1,108 @@
+use super::Locked;
+use core::alloc::{GlobalAlloc, Layout};
+
+/// Fixed size block allocator
+///
+/// Advantages of this allocator:
+///     - O(1) time complexity
+///
+/// Fallback: Linked list allocator
+
+/// Sizes are also used for alignments, so they need to be powers of 2.
+const BLOCK_SIZES: &[usize] = &[8, 16, 32, 64, 128, 256, 512, 1024, 2048];
+
+struct ListNode {
+    next: Option<&'static mut ListNode>,
+}
+
+pub struct FixedSizeBlockAllocator {
+    list_heads: [Option<&'static mut ListNode>; BLOCK_SIZES.len()],
+    fallback_allocator: linked_list_allocator::Heap,
+}
+
+impl FixedSizeBlockAllocator {
+    pub const fn new() -> Self {
+        const EMPTY: Option<&'static mut ListNode> = None;
+
+        FixedSizeBlockAllocator {
+            list_heads: [EMPTY; BLOCK_SIZES.len()],
+            // FIXME: Use the allocator that we created.
+            fallback_allocator: linked_list_allocator::Heap::empty(),
+        }
+    }
+
+    pub unsafe fn init(&mut self, heap_start: usize, heap_size: usize) {
+        unsafe {
+            self.fallback_allocator.init(heap_start, heap_size);
+        }
+    }
+
+    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_ptr(),
+        }
+    }
+}
+
+fn list_index(layout: &Layout) -> Option<usize> {
+    let required_block_size = layout.size().max(layout.align());
+
+    BLOCK_SIZES.iter().position(|&s| s >= required_block_size)
+}
+
+unsafe impl GlobalAlloc for Locked<FixedSizeBlockAllocator> {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let mut allocator = self.lock();
+
+        match 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 => {
+                    let block_size = BLOCK_SIZES[index];
+                    // `from_size_align` requires an alignment in powers of two.
+                    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 list_index(&layout) {
+            Some(index) => {
+                let new_node = ListNode {
+                    next: allocator.list_heads[index].take(),
+                };
+
+                assert!(mem::size_of::<ListNode>() <= BLOCK_SIZES[index]);
+                assert!(mem::size_of::<ListNode>() <= BLOCK_SIZES[index]);
+
+                let new_node_ptr = ptr as *mut ListNode;
+
+                unsafe {
+                    new_node_ptr.write(new_node);
+                    allocator.list_heads[index] = Some(&mut *new_node_ptr);
+                }
+            }
+
+            None => {
+                let ptr = NonNull::new(ptr).unwrap();
+
+                unsafe {
+                    allocator.fallback_allocator.deallocate(ptr, layout);
+                }
+            }
+        }
+    }
+}