Add a flag to prevent duplicate initialization

Author: JalonWong

CHANGELOG.md

@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 - Added a `init` macro to make initialization easier.
 
+### Changed
+
+- The `init` function will panic if it's called more than once or with `size == 0`.
+
 ## [v0.6.0] - 2024-09-01
 
 ### Added

src/lib.rs

@@ -30,6 +30,13 @@ pub use tlsf::Heap as TlsfHeap;
 /// It internally calls `Heap::init(...)` on the heap,
 /// so `Heap::init(...)` should not be called directly if this macro is used.
 ///
+/// # Panics
+///
+/// This macro will panic if either of the following are true:
+///
+/// - this function is called more than ONCE.
+/// - `size == 0`.
+///
 /// # Example
 ///
 /// ```rust

src/llff.rs

@@ -7,7 +7,7 @@ use linked_list_allocator::Heap as LLHeap;
 
 /// A linked list first fit heap.
 pub struct Heap {
-    heap: Mutex<RefCell<LLHeap>>,
+    heap: Mutex<RefCell<(LLHeap, bool)>>,
 }
 
 impl Heap {
@@ -17,7 +17,7 @@ impl Heap {
     /// [`init`](Self::init) method before using the allocator.
     pub const fn empty() -> Heap {
         Heap {
-            heap: Mutex::new(RefCell::new(LLHeap::empty())),
+            heap: Mutex::new(RefCell::new((LLHeap::empty(), false))),
         }
     }
 
@@ -41,34 +41,42 @@ impl Heap {
     ///
     /// # Safety
     ///
-    /// Obey these or Bad Stuff will happen.
+    /// This function is safe if the following invariants hold:
     ///
-    /// - This function must be called exactly ONCE.
-    /// - `size > 0`
+    /// - `start_addr` points to valid memory.
+    /// - `size` is correct.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if either of the following are true:
+    ///
+    /// - this function is called more than ONCE.
+    /// - `size == 0`.
     pub unsafe fn init(&self, start_addr: usize, size: usize) {
+        assert!(size > 0);
         critical_section::with(|cs| {
-            self.heap
-                .borrow(cs)
-                .borrow_mut()
-                .init(start_addr as *mut u8, size);
+            let mut heap = self.heap.borrow_ref_mut(cs);
+            assert!(!heap.1);
+            heap.1 = true;
+            heap.0.init(start_addr as *mut u8, size);
         });
     }
 
     /// Returns an estimate of the amount of bytes in use.
     pub fn used(&self) -> usize {
-        critical_section::with(|cs| self.heap.borrow(cs).borrow_mut().used())
+        critical_section::with(|cs| self.heap.borrow_ref_mut(cs).0.used())
     }
 
     /// Returns an estimate of the amount of bytes available.
     pub fn free(&self) -> usize {
-        critical_section::with(|cs| self.heap.borrow(cs).borrow_mut().free())
+        critical_section::with(|cs| self.heap.borrow_ref_mut(cs).0.free())
     }
 
     fn alloc(&self, layout: Layout) -> Option<NonNull<u8>> {
         critical_section::with(|cs| {
             self.heap
-                .borrow(cs)
-                .borrow_mut()
+                .borrow_ref_mut(cs)
+                .0
                 .allocate_first_fit(layout)
                 .ok()
         })
@@ -77,8 +85,8 @@ impl Heap {
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
         critical_section::with(|cs| {
             self.heap
-                .borrow(cs)
-                .borrow_mut()
+                .borrow_ref_mut(cs)
+                .0
                 .deallocate(NonNull::new_unchecked(ptr), layout)
         });
     }

src/tlsf.rs

@@ -10,7 +10,7 @@ type TlsfHeap = Tlsf<'static, usize, usize, { usize::BITS as usize }, { usize::B
 
 /// A two-Level segregated fit heap.
 pub struct Heap {
-    heap: Mutex<RefCell<TlsfHeap>>,
+    heap: Mutex<RefCell<(TlsfHeap, bool)>>,
 }
 
 impl Heap {
@@ -20,7 +20,7 @@ impl Heap {
     /// [`init`](Self::init) method before using the allocator.
     pub const fn empty() -> Heap {
         Heap {
-            heap: Mutex::new(RefCell::new(ConstDefault::DEFAULT)),
+            heap: Mutex::new(RefCell::new((ConstDefault::DEFAULT, false))),
         }
     }
 
@@ -44,29 +44,37 @@ impl Heap {
     ///
     /// # Safety
     ///
-    /// Obey these or Bad Stuff will happen.
+    /// This function is safe if the following invariants hold:
     ///
-    /// - This function must be called exactly ONCE.
-    /// - `size > 0`
+    /// - `start_addr` points to valid memory.
+    /// - `size` is correct.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if either of the following are true:
+    ///
+    /// - this function is called more than ONCE.
+    /// - `size == 0`.
     pub unsafe fn init(&self, start_addr: usize, size: usize) {
+        assert!(size > 0);
         critical_section::with(|cs| {
+            let mut heap = self.heap.borrow_ref_mut(cs);
+            assert!(!heap.1);
+            heap.1 = true;
             let block: &[u8] = core::slice::from_raw_parts(start_addr as *const u8, size);
-            self.heap
-                .borrow(cs)
-                .borrow_mut()
-                .insert_free_block_ptr(block.into());
+            heap.0.insert_free_block_ptr(block.into());
         });
     }
 
     fn alloc(&self, layout: Layout) -> Option<NonNull<u8>> {
-        critical_section::with(|cs| self.heap.borrow(cs).borrow_mut().allocate(layout))
+        critical_section::with(|cs| self.heap.borrow_ref_mut(cs).0.allocate(layout))
     }
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
         critical_section::with(|cs| {
             self.heap
-                .borrow(cs)
-                .borrow_mut()
+                .borrow_ref_mut(cs)
+                .0
                 .deallocate(NonNull::new_unchecked(ptr), layout.align())
         })
     }