Merge pull request #103 from JalonWong/master

Add a macro to make initialization easier

Author: zeenix

CHANGELOG.md

@@ -5,6 +5,16 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
 
+## [Unreleased]
+
+### Added
+
+- Added a `init` macro to make initialization easier.
+
+### Changed
+
+- The `Heap::init` methods now panic if they're called more than once or with `size == 0`.
+
 ## [v0.6.0] - 2024-09-01
 
 ### Added

Cargo.toml

@@ -29,7 +29,7 @@ allocator_api = []
 
 # Use the Two-Level Segregated Fit allocator
 tlsf = ["rlsf", "const-default"]
- # Use the LinkedList first-fit allocator
+# Use the LinkedList first-fit allocator
 llff = ["linked_list_allocator"]
 
 [dependencies]

README.md

@@ -31,6 +31,10 @@ static HEAP: Heap = Heap::empty();
 #[entry]
 fn main() -> ! {
     // Initialize the allocator BEFORE you use it
+    unsafe {
+        embedded_alloc::init!(HEAP, 1024);
+    }
+    // Alternatively, you can write the code directly to meet specific requirements.
     {
         use core::mem::MaybeUninit;
         const HEAP_SIZE: usize = 1024;

examples/global_alloc.rs

@@ -17,11 +17,8 @@ static HEAP: Heap = Heap::empty();
 #[entry]
 fn main() -> ! {
     // Initialize the allocator BEFORE you use it
-    {
-        use core::mem::MaybeUninit;
-        const HEAP_SIZE: usize = 1024;
-        static mut HEAP_MEM: [MaybeUninit<u8>; HEAP_SIZE] = [MaybeUninit::uninit(); HEAP_SIZE];
-        unsafe { HEAP.init(&raw mut HEAP_MEM as usize, HEAP_SIZE) }
+    unsafe {
+        embedded_alloc::init!(HEAP, 1024);
     }
 
     let mut xs = Vec::new();

examples/llff_integration_test.rs

@@ -63,10 +63,8 @@ fn test_allocator_api() {
 
 #[entry]
 fn main() -> ! {
-    {
-        const HEAP_SIZE: usize = 1024;
-        static mut HEAP_MEM: [MaybeUninit<u8>; HEAP_SIZE] = [MaybeUninit::uninit(); HEAP_SIZE];
-        unsafe { HEAP.init(&raw mut HEAP_MEM as usize, HEAP_SIZE) }
+    unsafe {
+        embedded_alloc::init!(HEAP, 1024);
     }
 
     #[allow(clippy::type_complexity)]

examples/tlsf_integration_test.rs

@@ -81,9 +81,8 @@ fn test_allocator_api() {
 
 #[entry]
 fn main() -> ! {
-    {
-        static mut HEAP_MEM: [MaybeUninit<u8>; HEAP_SIZE] = [MaybeUninit::uninit(); HEAP_SIZE];
-        unsafe { HEAP.init(&raw mut HEAP_MEM as usize, HEAP_SIZE) }
+    unsafe {
+        embedded_alloc::init!(HEAP, HEAP_SIZE);
     }
 
     #[allow(clippy::type_complexity)]

src/lib.rs

@@ -12,3 +12,55 @@ mod tlsf;
 pub use llff::Heap as LlffHeap;
 #[cfg(feature = "tlsf")]
 pub use tlsf::Heap as TlsfHeap;
+
+/// Initialize the global heap.
+///
+/// This macro creates a static, uninitialized memory buffer of the specified size and
+/// initializes the heap instance with that buffer.
+///
+/// # Parameters
+///
+/// - `$heap:ident`: The identifier of the global heap instance to initialize.
+/// - `$size:expr`: An expression evaluating to a `usize` that specifies the size of the
+///   static memory buffer in bytes. It must be **greater than zero**.
+///
+/// # Safety
+///
+/// This macro must be called first, before any operations on the heap, and **only once**.
+/// 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
+/// use cortex_m_rt::entry;
+/// use embedded_alloc::LlffHeap as Heap;
+///
+/// #[global_allocator]
+/// static HEAP: Heap = Heap::empty();
+///
+/// #[entry]
+/// fn main() -> ! {
+///     // Initialize the allocator BEFORE you use it
+///     unsafe {
+///         embedded_alloc::init!(HEAP, 1024);
+///     }
+///     let mut xs = Vec::new();
+///     // ...
+/// }
+/// ```
+#[macro_export]
+macro_rules! init {
+    ($heap:ident, $size:expr) => {
+        static mut HEAP_MEM: [::core::mem::MaybeUninit<u8>; $size] =
+            [::core::mem::MaybeUninit::uninit(); $size];
+        $heap.init(&raw mut HEAP_MEM as usize, $size)
+    };
+}

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())
         })
     }