From 3c6c7efc79bf8a8ac4bfbca6989fd24284ea140b Mon Sep 17 00:00:00 2001
From: Corwin <corwin@kuiper.dev>
Date: Mon, 7 Feb 2022 21:19:56 +0000
Subject: [PATCH] allocator that can have a custom start point

---
 agb/src/agb_alloc/block_allocator.rs |  4 ++--
 agb/src/agb_alloc/bump_allocator.rs  | 36 ++++------------------------
 agb/src/agb_alloc/mod.rs             | 27 ++++++++++++++++++++-
 3 files changed, 32 insertions(+), 35 deletions(-)

diff --git a/agb/src/agb_alloc/block_allocator.rs b/agb/src/agb_alloc/block_allocator.rs
index 6dfb154b..0866d697 100644
--- a/agb/src/agb_alloc/block_allocator.rs
+++ b/agb/src/agb_alloc/block_allocator.rs
@@ -49,9 +49,9 @@ pub(crate) struct BlockAllocator {
 }
 
 impl BlockAllocator {
-    pub(super) const unsafe fn new() -> Self {
+    pub(super) const unsafe fn new(start: fn() -> usize) -> Self {
         Self {
-            inner_allocator: BumpAllocator::new(),
+            inner_allocator: BumpAllocator::new(start),
             state: Mutex::new(RefCell::new(BlockAllocatorState {
                 first_free_block: None,
             })),
diff --git a/agb/src/agb_alloc/bump_allocator.rs b/agb/src/agb_alloc/bump_allocator.rs
index d1ddac3a..701f29df 100644
--- a/agb/src/agb_alloc/bump_allocator.rs
+++ b/agb/src/agb_alloc/bump_allocator.rs
@@ -8,12 +8,14 @@ use bare_metal::{CriticalSection, Mutex};
 
 pub(crate) struct BumpAllocator {
     current_ptr: Mutex<RefCell<Option<SendNonNull<u8>>>>,
+    start: Mutex<fn() -> usize>,
 }
 
 impl BumpAllocator {
-    pub const fn new() -> Self {
+    pub const fn new(start: fn() -> usize) -> Self {
         Self {
             current_ptr: Mutex::new(RefCell::new(None)),
+            start: Mutex::new(start),
         }
     }
 }
@@ -25,7 +27,7 @@ impl BumpAllocator {
         let ptr = if let Some(c) = *current_ptr {
             c.as_ptr() as usize
         } else {
-            get_data_end()
+            self.start.borrow(*cs)()
         };
 
         let alignment_bitmask = layout.align() - 1;
@@ -56,33 +58,3 @@ unsafe impl GlobalAlloc for BumpAllocator {
 
     unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
 }
-
-fn get_data_end() -> usize {
-    extern "C" {
-        static __ewram_data_end: usize;
-    }
-
-    // TODO: This seems completely wrong, but without the &, rust generates
-    // a double dereference :/. Maybe a bug in nightly?
-    (unsafe { &__ewram_data_end }) as *const _ as usize
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test_case]
-    fn should_return_data_end_somewhere_in_ewram(_gba: &mut crate::Gba) {
-        let data_end = get_data_end();
-
-        assert!(
-            0x0200_0000 <= data_end,
-            "data end should be bigger than 0x0200_0000, got {}",
-            data_end
-        );
-        assert!(
-            0x0204_0000 > data_end,
-            "data end should be smaller than 0x0203_0000"
-        );
-    }
-}
diff --git a/agb/src/agb_alloc/mod.rs b/agb/src/agb_alloc/mod.rs
index 4084f3af..c3f9a875 100644
--- a/agb/src/agb_alloc/mod.rs
+++ b/agb/src/agb_alloc/mod.rs
@@ -33,7 +33,7 @@ impl<T> DerefMut for SendNonNull<T> {
 const EWRAM_END: usize = 0x0204_0000;
 
 #[global_allocator]
-static GLOBAL_ALLOC: BlockAllocator = unsafe { BlockAllocator::new() };
+static GLOBAL_ALLOC: BlockAllocator = unsafe { BlockAllocator::new(get_data_end) };
 
 #[cfg(test)]
 pub unsafe fn number_of_blocks() -> u32 {
@@ -49,6 +49,16 @@ fn alloc_error(layout: Layout) -> ! {
     );
 }
 
+fn get_data_end() -> usize {
+    extern "C" {
+        static __ewram_data_end: usize;
+    }
+
+    // TODO: This seems completely wrong, but without the &, rust generates
+    // a double dereference :/. Maybe a bug in nightly?
+    (unsafe { &__ewram_data_end }) as *const _ as usize
+}
+
 #[cfg(test)]
 mod test {
     const EWRAM_START: usize = 0x0200_0000;
@@ -118,4 +128,19 @@ mod test {
         assert_eq!(v1[40], 137);
         assert_eq!(v2[78], 1075);
     }
+
+    #[test_case]
+    fn should_return_data_end_somewhere_in_ewram(_gba: &mut crate::Gba) {
+        let data_end = get_data_end();
+
+        assert!(
+            0x0200_0000 <= data_end,
+            "data end should be bigger than 0x0200_0000, got {}",
+            data_end
+        );
+        assert!(
+            0x0204_0000 > data_end,
+            "data end should be smaller than 0x0203_0000"
+        );
+    }
 }