Use OnceCell to prevent duplicate initialization

Author: JalonWong

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 ### Added
 
 - Added a `init` macro to make initialization easier.
+- Use `OnceCell` to prevent duplicate initialization.
 
 ## [v0.6.0] - 2024-09-01
 

src/lib.rs

@@ -30,6 +30,10 @@ 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
+///
+/// Calling this macro multiple times or with `size == 0` will cause a panic.
+///
 /// # Example
 ///
 /// ```rust

src/llff.rs

@@ -1,13 +1,13 @@
 use core::alloc::{GlobalAlloc, Layout};
-use core::cell::RefCell;
+use core::cell::{OnceCell, RefCell};
 use core::ptr::{self, NonNull};
 
 use critical_section::Mutex;
 use linked_list_allocator::Heap as LLHeap;
 
 /// A linked list first fit heap.
 pub struct Heap {
-    heap: Mutex<RefCell<LLHeap>>,
+    heap: Mutex<RefCell<OnceCell<LLHeap>>>,
 }
 
 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(OnceCell::new())),
         }
     }
 
@@ -41,34 +41,37 @@ impl Heap {
     ///
     /// # Safety
     ///
-    /// Obey these or Bad Stuff will happen.
+    /// This function will panic if either of the following are true:
     ///
-    /// - This function must be called exactly ONCE.
-    /// - `size > 0`
+    /// - 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);
+            assert!(self
+                .heap
+                .borrow_ref_mut(cs)
+                .set(LLHeap::new(start_addr as *mut u8, size))
+                .is_ok());
         });
     }
 
     /// 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).get_mut().unwrap().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).get_mut().unwrap().free())
     }
 
     fn alloc(&self, layout: Layout) -> Option<NonNull<u8>> {
         critical_section::with(|cs| {
             self.heap
-                .borrow(cs)
-                .borrow_mut()
+                .borrow_ref_mut(cs)
+                .get_mut()
+                .unwrap()
                 .allocate_first_fit(layout)
                 .ok()
         })
@@ -77,8 +80,9 @@ 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)
+                .get_mut()
+                .unwrap()
                 .deallocate(NonNull::new_unchecked(ptr), layout)
         });
     }

src/tlsf.rs

@@ -1,16 +1,15 @@
 use core::alloc::{GlobalAlloc, Layout};
-use core::cell::RefCell;
+use core::cell::{OnceCell, RefCell};
 use core::ptr::{self, NonNull};
 
-use const_default::ConstDefault;
 use critical_section::Mutex;
 use rlsf::Tlsf;
 
 type TlsfHeap = Tlsf<'static, usize, usize, { usize::BITS as usize }, { usize::BITS as usize }>;
 
 /// A two-Level segregated fit heap.
 pub struct Heap {
-    heap: Mutex<RefCell<TlsfHeap>>,
+    heap: Mutex<RefCell<OnceCell<TlsfHeap>>>,
 }
 
 impl Heap {
@@ -20,7 +19,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(OnceCell::new())),
         }
     }
 
@@ -44,29 +43,39 @@ impl Heap {
     ///
     /// # Safety
     ///
-    /// Obey these or Bad Stuff will happen.
+    /// This function will panic if either of the following are true:
     ///
-    /// - This function must be called exactly ONCE.
-    /// - `size > 0`
+    /// - 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 block: &[u8] = core::slice::from_raw_parts(start_addr as *const u8, size);
+            self.heap.borrow_ref_mut(cs).set(TlsfHeap::new()).unwrap();
             self.heap
-                .borrow(cs)
-                .borrow_mut()
+                .borrow_ref_mut(cs)
+                .get_mut()
+                .unwrap()
                 .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)
+                .get_mut()
+                .unwrap()
+                .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)
+                .get_mut()
+                .unwrap()
                 .deallocate(NonNull::new_unchecked(ptr), layout.align())
         })
     }