create linked list allocator

Author: ash bek

src/allocator/linked_list.rs

@@ -0,0 +1,159 @@
+use super::{align_up, Locked};
+use alloc::alloc::{GlobalAlloc, Layout};
+use core::{mem, ptr};
+
+/// List node for a linked list allocator
+///
+/// The node contains its size and a pointer to the next
+/// node.
+struct ListNode {
+    size: usize,
+    next: Option<&'static mut ListNode>,
+}
+
+impl ListNode {
+    const fn new(size: usize) -> Self {
+        ListNode { size, next: None }
+    }
+
+    /// Return the start address of the node.
+    fn start_addr(&self) -> usize {
+        self as *const Self as usize
+    }
+
+    /// Return the end address of the region
+    ///
+    /// This works because the region itself contains the node, so the end address
+    /// should always be `self.start_addr() + self.size`
+    fn end_addr(&self) -> usize {
+        self.start_addr() + self.size
+    }
+}
+
+pub struct LinkedListAllocator {
+    head: ListNode,
+}
+
+impl LinkedListAllocator {
+    /// Create a new empty linked list allocator
+    pub const fn new() -> Self {
+        LinkedListAllocator {
+            head: ListNode::new(0),
+        }
+    }
+
+    /// Initialize the allocator with the given heap bounds.
+    ///
+    /// This function is unsafe because the caller must guarantee that the given
+    /// heap bounds are valid and that the heap is unused. This method must be
+    /// called only once.
+    pub unsafe fn init(&mut self, heap_start: usize, heap_size: usize) {
+        self.add_free_region(heap_start, heap_size)
+    }
+
+    /// Add a new free region to the allocator
+    ///
+    /// This function is unsafe as we're operating on raw memory and need to ensure
+    /// the alignment and size is correct.
+    unsafe fn add_free_region(&mut self, addr: usize, size: usize) {
+        // Ensure the size of the free region is capable of holding a node
+        assert_eq!(align_up(addr, mem::align_of::<ListNode>()), addr);
+        assert!(size >= mem::size_of::<ListNode>());
+
+        // 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)
+    }
+
+    /// Looks for a free region with the given size and alignment and removes
+    /// it from the list.
+    ///
+    /// Returns a tuple of the list node and the start address of the allocation.
+    fn find_region(&mut self, size: usize, align: usize) -> Option<(&'static mut ListNode, usize)> {
+        // Ref to the current list node, updated each iteration.
+        let mut current = &mut self.head;
+
+        // Iterate over the list until we find a large enough region.
+        while let Some(ref mut region) = current.next {
+            if let Ok(alloc_start) = Self::alloc_from_region(&region, size, align) {
+                // Region suitable for allocation, remove it from the list.
+                let next = region.next.take();
+                let ret = Some((current.next.take().unwrap(), alloc_start));
+                current.next = next;
+                return ret;
+            } else {
+                // Region not suitable, find another region.
+                current = current.next.as_mut().unwrap();
+            }
+        }
+        // No suitable regions.
+        None
+    }
+
+    /// Try to use the given region for an allocation with given size and
+    /// alignment.
+    ///
+    /// Returns the allocation start address on success.
+    fn alloc_from_region(region: &ListNode, size: usize, align: usize) -> Result<usize, ()> {
+        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 is too small for the desired allocation
+            return Err(());
+        }
+
+        // Region is large enough for the required allocation, now check if the leftover region can contain a node.
+        let excess_size = region.end_addr() - alloc_end;
+        if excess_size > 0 && excess_size < mem::size_of::<ListNode>() {
+            return Err(());
+        }
+
+        // Region suitable for allocation.
+        Ok(alloc_start)
+    }
+
+    /// Adjust the given layout so that the resulting allocated memory
+    /// region is also capable of storing a `ListNode`.
+    ///
+    /// Returns the adjusted size and alignment as a (size, align) tuple.
+    fn size_align(layout: Layout) -> (usize, usize) {
+        let layout = layout
+            .align_to(mem::align_of::<ListNode>())
+            .expect("Alignment failed")
+            .pad_to_align();
+
+        let size = layout.size().max(mem::size_of::<ListNode>());
+
+        (size, layout.align())
+    }
+}
+
+unsafe impl GlobalAlloc for Locked<LinkedListAllocator> {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        // Perform normal 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("Size has overflowed");
+            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) {
+        let (size, _) = LinkedListAllocator::size_align(layout);
+        self.lock().add_free_region(ptr as usize, size)
+    }
+}