From 03e95172157e93b4430ba861b7e0a65fa854ac16 Mon Sep 17 00:00:00 2001
From: Corwin Kuiper <corwin@kuiper.dev>
Date: Mon, 17 Jan 2022 21:46:14 +0000
Subject: [PATCH] only use critical section once

---
 agb/src/agb_alloc/block_allocator.rs |  8 ++---
 agb/src/agb_alloc/bump_allocator.rs  | 53 ++++++++++++++--------------
 2 files changed, 31 insertions(+), 30 deletions(-)

diff --git a/agb/src/agb_alloc/block_allocator.rs b/agb/src/agb_alloc/block_allocator.rs
index 1fc17a1b..f589de46 100644
--- a/agb/src/agb_alloc/block_allocator.rs
+++ b/agb/src/agb_alloc/block_allocator.rs
@@ -4,7 +4,7 @@ use core::cell::RefCell;
 use core::ptr::NonNull;
 
 use crate::interrupt::free;
-use bare_metal::Mutex;
+use bare_metal::{CriticalSection, Mutex};
 
 use super::bump_allocator::BumpAllocator;
 use super::SendNonNull;
@@ -47,9 +47,9 @@ impl BlockAllocator {
         }
     }
 
-    fn new_block(&self, layout: Layout) -> *mut u8 {
+    fn new_block(&self, layout: Layout, cs: &CriticalSection) -> *mut u8 {
         let overall_layout = Block::either_layout(layout);
-        self.inner_allocator.alloc_safe(overall_layout)
+        self.inner_allocator.alloc_critical(overall_layout, cs)
     }
 }
 
@@ -91,7 +91,7 @@ unsafe impl GlobalAlloc for BlockAllocator {
                     list_ptr = &mut curr_block.next;
                 }
 
-                self.new_block(layout)
+                self.new_block(layout, key)
             })
         }
     }
diff --git a/agb/src/agb_alloc/bump_allocator.rs b/agb/src/agb_alloc/bump_allocator.rs
index 671d4abe..b5e9b90a 100644
--- a/agb/src/agb_alloc/bump_allocator.rs
+++ b/agb/src/agb_alloc/bump_allocator.rs
@@ -4,7 +4,7 @@ use core::ptr::NonNull;
 
 use super::SendNonNull;
 use crate::interrupt::free;
-use bare_metal::Mutex;
+use bare_metal::{CriticalSection, Mutex};
 
 pub(crate) struct BumpAllocator {
     current_ptr: Mutex<RefCell<Option<SendNonNull<u8>>>>,
@@ -19,32 +19,33 @@ impl BumpAllocator {
 }
 
 impl BumpAllocator {
+    pub fn alloc_critical(&self, layout: Layout, cs: &CriticalSection) -> *mut u8 {
+        let mut current_ptr = self.current_ptr.borrow(*cs).borrow_mut();
+
+        let ptr = if let Some(c) = *current_ptr {
+            c.as_ptr() as usize
+        } else {
+            get_data_end()
+        };
+
+        let alignment_bitmask = layout.align() - 1;
+        let fixup = ptr & alignment_bitmask;
+
+        let amount_to_add = layout.align() - fixup;
+
+        let resulting_ptr = ptr + amount_to_add;
+        let new_current_ptr = resulting_ptr + layout.size();
+
+        if new_current_ptr as usize >= super::EWRAM_END {
+            return core::ptr::null_mut();
+        }
+
+        *current_ptr = NonNull::new(new_current_ptr as *mut _).map(SendNonNull);
+
+        resulting_ptr as *mut _
+    }
     pub fn alloc_safe(&self, layout: Layout) -> *mut u8 {
-        free(|key| {
-            let mut current_ptr = self.current_ptr.borrow(*key).borrow_mut();
-
-            let ptr = if let Some(c) = *current_ptr {
-                c.as_ptr() as usize
-            } else {
-                get_data_end()
-            };
-
-            let alignment_bitmask = layout.align() - 1;
-            let fixup = ptr & alignment_bitmask;
-
-            let amount_to_add = layout.align() - fixup;
-
-            let resulting_ptr = ptr + amount_to_add;
-            let new_current_ptr = resulting_ptr + layout.size();
-
-            if new_current_ptr as usize >= super::EWRAM_END {
-                return core::ptr::null_mut();
-            }
-
-            *current_ptr = NonNull::new(new_current_ptr as *mut _).map(SendNonNull);
-
-            resulting_ptr as *mut _
-        })
+        free(|key| self.alloc_critical(layout, key))
     }
 }