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,54 +106,77 @@ 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) =
+        let (block_after_layout, block_after_layout_offset) = full_layout
-            full_layout.extend(Layout::new::<Block>()).unwrap();
+            .extend(Layout::new::<Block>().align_to(8).unwrap().pad_to_align())
+            .unwrap();
 
-        {
+        free(|key| {
-            free(|key| {
+            let mut state = self.state.borrow(*key).borrow_mut();
-                let mut state = self.state.borrow(*key).borrow_mut();
+            let mut current_block = state.first_free_block;
-                let mut current_block = state.first_free_block;
+            let mut list_ptr = &mut state.first_free_block;
-                let mut list_ptr = &mut state.first_free_block;
+            while let Some(mut curr) = current_block {
-                while let Some(mut curr) = current_block {
+                let curr_block = curr.as_mut();
-                    let curr_block = curr.as_mut();
+                if curr_block.size == full_layout.size() {
-                    if curr_block.size == full_layout.size() {
+                    *list_ptr = curr_block.next;
-                        *list_ptr = curr_block.next;
+                    return curr.as_ptr().cast();
-                        return curr.as_ptr().cast();
+                } else if curr_block.size >= block_after_layout.size() {
-                    } else if curr_block.size >= block_after_layout.size() {
+                    // can split block
-                        // can split block
+                    let split_block = Block {
-                        let split_block = Block {
+                        size: curr_block.size - block_after_layout_offset,
-                            size: curr_block.size - block_after_layout_offset,
+                        next: curr_block.next,
-                            next: curr_block.next,
+                    };
-                        };
+                    let split_ptr = curr
-                        let split_ptr = curr
+                        .as_ptr()
-                            .as_ptr()
+                        .cast::<u8>()
-                            .cast::<u8>()
+                        .add(block_after_layout_offset)
-                            .add(block_after_layout_offset)
+                        .cast();
-                            .cast();
+                    *split_ptr = split_block;
-                        *split_ptr = split_block;
+                    *list_ptr = NonNull::new(split_ptr).map(SendNonNull);
-                        *list_ptr = NonNull::new(split_ptr).map(SendNonNull);
 
-                        return curr.as_ptr().cast();
+                    return curr.as_ptr().cast();
-                    }
-                    current_block = curr_block.next;
-                    list_ptr = &mut curr_block.next;
                 }
+                current_block = curr_block.next;
+                list_ptr = &mut curr_block.next;
+            }
 
-                self.new_block(layout, key)
+            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 new_block_content = Block {
+
-                size: new_layout.size(),
+            let mut list_ptr = &mut state.first_free_block;
-                next: state.first_free_block,
+
-            };
+            loop {
-            *ptr.cast() = new_block_content;
+                match list_ptr {
-            state.first_free_block = NonNull::new(ptr.cast()).map(SendNonNull);
+                    Some(mut current_block) => {
-        })
+                        if current_block.as_ptr().cast() > ptr {
+                            let new_block_content = Block {
+                                size: new_layout.size(),
+                                next: Some(current_block),
+                            };
+                            *ptr.cast() = new_block_content;
+                            *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!(