From 475edadcb92841e7382a1a952536d2fd14adfa96 Mon Sep 17 00:00:00 2001
From: Corwin <corwin@kuiper.dev>
Date: Mon, 7 Feb 2022 21:43:17 +0000
Subject: [PATCH] extend to also supporting the end

---
 agb/src/agb_alloc/block_allocator.rs |  4 ++--
 agb/src/agb_alloc/bump_allocator.rs  | 17 ++++++++++++-----
 agb/src/agb_alloc/mod.rs             |  9 ++++++++-
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/agb/src/agb_alloc/block_allocator.rs b/agb/src/agb_alloc/block_allocator.rs
index 0866d69..1d2221a 100644
--- a/agb/src/agb_alloc/block_allocator.rs
+++ b/agb/src/agb_alloc/block_allocator.rs
@@ -12,7 +12,7 @@ use core::ptr::NonNull;
 use crate::interrupt::free;
 use bare_metal::{CriticalSection, Mutex};
 
-use super::bump_allocator::BumpAllocator;
+use super::bump_allocator::{BumpAllocator, StartEnd};
 use super::SendNonNull;
 
 struct Block {
@@ -49,7 +49,7 @@ pub(crate) struct BlockAllocator {
 }
 
 impl BlockAllocator {
-    pub(super) const unsafe fn new(start: fn() -> usize) -> Self {
+    pub(super) const unsafe fn new(start: StartEnd) -> Self {
         Self {
             inner_allocator: BumpAllocator::new(start),
             state: Mutex::new(RefCell::new(BlockAllocatorState {
diff --git a/agb/src/agb_alloc/bump_allocator.rs b/agb/src/agb_alloc/bump_allocator.rs
index 701f29d..5b6ca15 100644
--- a/agb/src/agb_alloc/bump_allocator.rs
+++ b/agb/src/agb_alloc/bump_allocator.rs
@@ -6,16 +6,23 @@ use super::SendNonNull;
 use crate::interrupt::free;
 use bare_metal::{CriticalSection, Mutex};
 
+pub(crate) struct AddrFn(pub fn() -> usize);
+
+pub(crate) struct StartEnd {
+    pub start: AddrFn,
+    pub end: AddrFn,
+}
+
 pub(crate) struct BumpAllocator {
     current_ptr: Mutex<RefCell<Option<SendNonNull<u8>>>>,
-    start: Mutex<fn() -> usize>,
+    start_end: Mutex<StartEnd>,
 }
 
 impl BumpAllocator {
-    pub const fn new(start: fn() -> usize) -> Self {
+    pub const fn new(start_end: StartEnd) -> Self {
         Self {
             current_ptr: Mutex::new(RefCell::new(None)),
-            start: Mutex::new(start),
+            start_end: Mutex::new(start_end),
         }
     }
 }
@@ -27,7 +34,7 @@ impl BumpAllocator {
         let ptr = if let Some(c) = *current_ptr {
             c.as_ptr() as usize
         } else {
-            self.start.borrow(*cs)()
+            self.start_end.borrow(*cs).start.0()
         };
 
         let alignment_bitmask = layout.align() - 1;
@@ -38,7 +45,7 @@ impl BumpAllocator {
         let resulting_ptr = ptr + amount_to_add;
         let new_current_ptr = resulting_ptr + layout.size();
 
-        if new_current_ptr as usize >= super::EWRAM_END {
+        if new_current_ptr as usize >= self.start_end.borrow(*cs).end.0() {
             return core::ptr::null_mut();
         }
 
diff --git a/agb/src/agb_alloc/mod.rs b/agb/src/agb_alloc/mod.rs
index c3f9a87..97fc173 100644
--- a/agb/src/agb_alloc/mod.rs
+++ b/agb/src/agb_alloc/mod.rs
@@ -7,6 +7,8 @@ mod bump_allocator;
 
 use block_allocator::BlockAllocator;
 
+use self::bump_allocator::{AddrFn, StartEnd};
+
 struct SendNonNull<T>(NonNull<T>);
 unsafe impl<T> Send for SendNonNull<T> {}
 
@@ -33,7 +35,12 @@ impl<T> DerefMut for SendNonNull<T> {
 const EWRAM_END: usize = 0x0204_0000;
 
 #[global_allocator]
-static GLOBAL_ALLOC: BlockAllocator = unsafe { BlockAllocator::new(get_data_end) };
+static GLOBAL_ALLOC: BlockAllocator = unsafe {
+    BlockAllocator::new(StartEnd {
+        start: AddrFn(get_data_end),
+        end: AddrFn(|| EWRAM_END),
+    })
+};
 
 #[cfg(test)]
 pub unsafe fn number_of_blocks() -> u32 {