A nice allocator!! Merges blocks together

agb/src/agb_alloc/block_allocator.rs

@@ -1,6 +1,7 @@
 use core::alloc::{GlobalAlloc, Layout};
 
 use core::cell::RefCell;
+use core::convert::TryInto;
 use core::ptr::NonNull;
 
 use crate::interrupt::free;
@@ -25,6 +26,8 @@ impl Block {
             aligned_to.align(),
         )
         .expect("too large allocation")
+        .align_to(8)
+        .expect("too large allocation")
     }
 }
 
@@ -47,10 +50,55 @@ impl BlockAllocator {
         }
     }
 
+    pub unsafe fn number_of_blocks(&self) -> u32 {
+        free(|key| {
+            let mut state = self.state.borrow(*key).borrow_mut();
+
+            let mut count = 0;
+
+            let mut list_ptr = &mut state.first_free_block;
+            while let Some(mut curr) = list_ptr {
+                count += 1;
+                list_ptr = &mut curr.as_mut().next;
+            }
+
+            count
+        })
+    }
+
     fn new_block(&self, layout: Layout, cs: &CriticalSection) -> *mut u8 {
         let overall_layout = Block::either_layout(layout);
         self.inner_allocator.alloc_critical(overall_layout, cs)
     }
+
+    unsafe fn normalise(&self) {
+        free(|key| {
+            let mut state = self.state.borrow(*key).borrow_mut();
+
+            let mut list_ptr = &mut state.first_free_block;
+
+            while let Some(mut curr) = list_ptr {
+                if let Some(next_elem) = curr.as_mut().next {
+                    let difference = next_elem
+                        .as_ptr()
+                        .cast::<u8>()
+                        .offset_from(curr.as_ptr().cast::<u8>());
+                    let usize_difference: usize = difference
+                        .try_into()
+                        .expect("distances in alloc'd blocks must be positive");
+
+                    if usize_difference == curr.as_mut().size {
+                        let current = curr.as_mut();
+                        let next = next_elem.as_ref();
+
+                        current.size += next.size;
+                        current.next = next.next;
+                    }
+                }
+                list_ptr = &mut curr.as_mut().next;
+            }
+        });
+    }
 }
 
 unsafe impl GlobalAlloc for BlockAllocator {
@@ -58,10 +106,10 @@ unsafe impl GlobalAlloc for BlockAllocator {
         // find a block that this current request fits in
         let full_layout = Block::either_layout(layout);
 
-        let (block_after_layout, block_after_layout_offset) =
-            full_layout.extend(Layout::new::<Block>()).unwrap();
+        let (block_after_layout, block_after_layout_offset) = full_layout
+            .extend(Layout::new::<Block>().align_to(8).unwrap().pad_to_align())
+            .unwrap();
 
-        {
         free(|key| {
             let mut state = self.state.borrow(*key).borrow_mut();
             let mut current_block = state.first_free_block;
@@ -94,18 +142,41 @@ unsafe impl GlobalAlloc for BlockAllocator {
             self.new_block(layout, key)
         })
     }
-    }
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-        let new_layout = Block::either_layout(layout);
+        let new_layout = Block::either_layout(layout).pad_to_align();
         free(|key| {
             let mut state = self.state.borrow(*key).borrow_mut();
+
+            let mut list_ptr = &mut state.first_free_block;
+
+            loop {
+                match list_ptr {
+                    Some(mut current_block) => {
+                        if current_block.as_ptr().cast() > ptr {
                             let new_block_content = Block {
                                 size: new_layout.size(),
-                next: state.first_free_block,
+                                next: Some(current_block),
                             };
                             *ptr.cast() = new_block_content;
-            state.first_free_block = NonNull::new(ptr.cast()).map(SendNonNull);
-        })
+                            *list_ptr = NonNull::new(ptr.cast()).map(SendNonNull);
+                            break;
+                        }
+                        list_ptr = &mut current_block.as_mut().next;
+                    }
+                    None => {
+                        // reached the end of the list without finding a place to insert the value
+                        let new_block_content = Block {
+                            size: new_layout.size(),
+                            next: None,
+                        };
+                        *ptr.cast() = new_block_content;
+                        *list_ptr = NonNull::new(ptr.cast()).map(SendNonNull);
+                        break;
+                    }
+                }
+            }
+        });
+        self.normalise();
     }
 }

agb/src/agb_alloc/bump_allocator.rs

@@ -31,7 +31,7 @@ impl BumpAllocator {
         let alignment_bitmask = layout.align() - 1;
         let fixup = ptr & alignment_bitmask;
 
-        let amount_to_add = layout.align() - fixup;
+        let amount_to_add = (layout.align() - fixup) & alignment_bitmask;
 
         let resulting_ptr = ptr + amount_to_add;
         let new_current_ptr = resulting_ptr + layout.size();

agb/src/agb_alloc/mod.rs

@@ -35,6 +35,10 @@ const EWRAM_END: usize = 0x0204_0000;
 #[global_allocator]
 static GLOBAL_ALLOC: BlockAllocator = unsafe { BlockAllocator::new() };
 
+pub unsafe fn number_of_blocks() -> u32 {
+    GLOBAL_ALLOC.number_of_blocks()
+}
+
 #[alloc_error_handler]
 fn alloc_error(layout: Layout) -> ! {
     panic!(