esp-wifi uses global allocator, esp-alloc supports multiple regions

esp-alloc/CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- a global allocator is created in esp-alloc, now you need to add individual memory regions (up to 3) to the allocator (#2099)
+
 ### Fixed
 
 ### Removed

esp-alloc/Cargo.toml

@@ -23,7 +23,8 @@ default-target = "riscv32imc-unknown-none-elf"
 features       = ["nightly"]
 
 [dependencies]
-critical-section      = "1.1.2"
+critical-section      = "1.1.3"
+enumset               = "1.1.5"
 linked_list_allocator = { version = "0.10.5", default-features = false, features = ["const_mut_refs"] }
 
 [features]

esp-alloc/src/lib.rs

@@ -1,33 +1,35 @@
-//! A simple `no_std` heap allocator for RISC-V and Xtensa processors from
+//! A `no_std` heap allocator for RISC-V and Xtensa processors from
 //! Espressif. Supports all currently available ESP32 devices.
 //!
 //! **NOTE:** using this as your global allocator requires using Rust 1.68 or
 //! greater, or the `nightly` release channel.
 //!
 //! # Using this as your Global Allocator
-//! To use EspHeap as your global allocator, you need at least Rust 1.68 or
-//! nightly.
 //!
 //! ```rust
-//! #[global_allocator]
-//! static ALLOCATOR: esp_alloc::EspHeap = esp_alloc::EspHeap::empty();
+//! use esp_alloc as _;
 //!
 //! fn init_heap() {
 //!     const HEAP_SIZE: usize = 32 * 1024;
 //!     static mut HEAP: MaybeUninit<[u8; HEAP_SIZE]> = MaybeUninit::uninit();
 //!
 //!     unsafe {
-//!         ALLOCATOR.init(HEAP.as_mut_ptr() as *mut u8, HEAP_SIZE);
+//!         esp_alloc::INSTANCE.add_region(esp_alloc::HeapRegion::new(
+//!             HEAP.as_mut_ptr() as *mut u8,
+//!             HEAP_SIZE,
+//!             esp_alloc::MemoryCapability::Internal.into(),
+//!         ));
 //!     }
 //! }
 //! ```
 //!
 //! # Using this with the nightly `allocator_api`-feature
-//! Sometimes you want to have single allocations in PSRAM, instead of an esp's
-//! DRAM. For that, it's convenient to use the nightly `allocator_api`-feature,
+//! Sometimes you want to have more control over allocations.
+//!
+//! For that, it's convenient to use the nightly `allocator_api`-feature,
 //! which allows you to specify an allocator for single allocations.
 //!
-//! **NOTE:** To use this, you have to enable the create's `nightly` feature
+//! **NOTE:** To use this, you have to enable the crate's `nightly` feature
 //! flag.
 //!
 //! Create and initialize an allocator to use in single allocations:
@@ -36,7 +38,11 @@
 //!
 //! fn init_psram_heap() {
 //!     unsafe {
-//!         PSRAM_ALLOCATOR.init(psram::psram_vaddr_start() as *mut u8, psram::PSRAM_BYTES);
+//!         PSRAM_ALLOCATOR.add_region(esp_alloc::HeapRegion::new(
+//!             psram::psram_vaddr_start() as *mut u8,
+//!             psram::PSRAM_BYTES,
+//!             esp_alloc::MemoryCapability::Internal.into(),
+//!         ));
 //!     }
 //! }
 //! ```
@@ -50,7 +56,7 @@
 #![cfg_attr(feature = "nightly", feature(allocator_api))]
 #![doc(html_logo_url = "https://avatars.githubusercontent.com/u/46717278")]
 
-pub mod macros;
+mod macros;
 
 #[cfg(feature = "nightly")]
 use core::alloc::{AllocError, Allocator};
@@ -61,35 +67,81 @@ use core::{
 };
 
 use critical_section::Mutex;
+use enumset::{EnumSet, EnumSetType};
 use linked_list_allocator::Heap;
 
+/// The global allocator instance
+#[global_allocator]
+pub static HEAP: EspHeap = EspHeap::empty();
+
+const NON_REGION: Option<HeapRegion> = None;
+
+#[derive(EnumSetType)]
+/// Describes the properties of a memory region
+pub enum MemoryCapability {
+    /// Memory must be internal; specifically it should not disappear when
+    /// flash/spiram cache is switched off
+    Internal,
+    /// Memory must be in SPI RAM
+    External,
+}
+
+/// A memory region to be used as heap memory
+pub struct HeapRegion {
+    heap: Heap,
+    capabilities: EnumSet<MemoryCapability>,
+}
+
+impl HeapRegion {
+    /// Create a new [HeapRegion] with the given capabilities
+    ///
+    /// # Safety
+    ///
+    /// - The supplied memory region must be available for the entire program
+    ///   (`'static`).
+    /// - The supplied memory region must be exclusively available to the heap
+    ///   only, no aliasing.
+    /// - `size > 0`.
+    pub unsafe fn new(
+        heap_bottom: *mut u8,
+        size: usize,
+        capabilities: EnumSet<MemoryCapability>,
+    ) -> Self {
+        let mut heap = Heap::empty();
+        heap.init(heap_bottom, size);
+
+        Self { heap, capabilities }
+    }
+}
+
+/// A memory allocator
+///
+/// In addition to what Rust's memory allocator can do it allows to allocate
+/// memory in regions satisfying specific needs.
 pub struct EspHeap {
-    heap: Mutex<RefCell<Heap>>,
+    heap: Mutex<RefCell<[Option<HeapRegion>; 3]>>,
 }
 
 impl EspHeap {
     /// Crate a new UNINITIALIZED heap allocator
-    ///
-    /// You must initialize this heap using the
-    /// [`init`](struct.EspHeap.html#method.init) method before using the
-    /// allocator.
-    pub const fn empty() -> EspHeap {
+    pub const fn empty() -> Self {
         EspHeap {
-            heap: Mutex::new(RefCell::new(Heap::empty())),
+            heap: Mutex::new(RefCell::new([NON_REGION; 3])),
         }
     }
 
-    /// Initializes the heap
-    ///
-    /// This function must be called BEFORE you run any code that makes use of
-    /// the allocator.
+    /// Add a memory region to the heap
     ///
     /// `heap_bottom` is a pointer to the location of the bottom of the heap.
     ///
     /// `size` is the size of the heap in bytes.
     ///
+    /// You can add up to three regions per allocator.
+    ///
     /// Note that:
     ///
+    /// - Memory is allocated from the first suitable memory region first
+    ///
     /// - The heap grows "upwards", towards larger addresses. Thus `end_addr`
     ///   must be larger than `start_addr`
     ///
@@ -102,59 +154,147 @@ impl EspHeap {
     ///   `'static` lifetime).
     /// - The supplied memory region must be exclusively available to the heap
     ///   only, no aliasing.
-    /// - This function must be called exactly ONCE.
     /// - `size > 0`.
-    pub unsafe fn init(&self, heap_bottom: *mut u8, size: usize) {
-        critical_section::with(|cs| self.heap.borrow(cs).borrow_mut().init(heap_bottom, size));
+    pub unsafe fn add_region(&self, region: HeapRegion) {
+        critical_section::with(|cs| {
+            let mut regions = self.heap.borrow_ref_mut(cs);
+            let free = regions
+                .iter()
+                .enumerate()
+                .find(|v| v.1.is_none())
+                .map(|v| v.0);
+
+            if let Some(free) = free {
+                regions[free] = Some(region);
+            } else {
+                panic!(
+                    "Exceeded the maximum of {} heap memory regions",
+                    regions.len()
+                );
+            }
+        });
     }
 
-    /// Returns an estimate of the amount of bytes in use.
+    /// Returns an estimate of the amount of bytes in use in all memory regions.
     pub fn used(&self) -> usize {
-        critical_section::with(|cs| self.heap.borrow(cs).borrow_mut().used())
+        critical_section::with(|cs| {
+            let regions = self.heap.borrow_ref(cs);
+            let mut used = 0;
+            for region in regions.iter() {
+                if let Some(region) = region.as_ref() {
+                    used += region.heap.used();
+                }
+            }
+            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())
+        self.free_caps(EnumSet::empty())
+    }
+
+    /// The free heap satisfying the given requirements
+    pub fn free_caps(&self, capabilities: EnumSet<MemoryCapability>) -> usize {
+        critical_section::with(|cs| {
+            let regions = self.heap.borrow_ref(cs);
+            let mut free = 0;
+            for region in regions.iter().filter(|region| {
+                if region.is_some() {
+                    region
+                        .as_ref()
+                        .unwrap()
+                        .capabilities
+                        .is_superset(capabilities)
+                } else {
+                    false
+                }
+            }) {
+                if let Some(region) = region.as_ref() {
+                    free += region.heap.free();
+                }
+            }
+            free
+        })
+    }
+
+    /// Allocate memory in a region satisfying the given requirements.
+    ///
+    /// # Safety
+    ///
+    /// This function is unsafe because undefined behavior can result
+    /// if the caller does not ensure that `layout` has non-zero size.
+    ///
+    /// The allocated block of memory may or may not be initialized.
+    pub unsafe fn alloc_caps(
+        &self,
+        capabilities: EnumSet<MemoryCapability>,
+        layout: Layout,
+    ) -> *mut u8 {
+        critical_section::with(|cs| {
+            let mut regions = self.heap.borrow_ref_mut(cs);
+            let mut iter = (*regions).iter_mut().filter(|region| {
+                if region.is_some() {
+                    region
+                        .as_ref()
+                        .unwrap()
+                        .capabilities
+                        .is_superset(capabilities)
+                } else {
+                    false
+                }
+            });
+
+            let res = loop {
+                if let Some(Some(region)) = iter.next() {
+                    let res = region.heap.allocate_first_fit(layout);
+                    if let Ok(res) = res {
+                        break Some(res);
+                    }
+                } else {
+                    break None;
+                }
+            };
+
+            res.map_or(ptr::null_mut(), |allocation| allocation.as_ptr())
+        })
     }
 }
 
 unsafe impl GlobalAlloc for EspHeap {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        critical_section::with(|cs| {
-            self.heap
-                .borrow(cs)
-                .borrow_mut()
-                .allocate_first_fit(layout)
-                .ok()
-                .map_or(ptr::null_mut(), |allocation| allocation.as_ptr())
-        })
+        self.alloc_caps(EnumSet::empty(), layout)
     }
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        if ptr.is_null() {
+            return;
+        }
+
         critical_section::with(|cs| {
-            self.heap
-                .borrow(cs)
-                .borrow_mut()
-                .deallocate(NonNull::new_unchecked(ptr), layout)
-        });
+            let mut regions = self.heap.borrow_ref_mut(cs);
+            let mut iter = (*regions).iter_mut();
+
+            while let Some(Some(region)) = iter.next() {
+                if region.heap.bottom() <= ptr && region.heap.top() >= ptr {
+                    region.heap.deallocate(NonNull::new_unchecked(ptr), layout);
+                }
+            }
+        })
     }
 }
 
 #[cfg(feature = "nightly")]
 unsafe impl Allocator for EspHeap {
     fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
-        critical_section::with(|cs| {
-            let raw_ptr = self
-                .heap
-                .borrow(cs)
-                .borrow_mut()
-                .allocate_first_fit(layout)
-                .map_err(|_| AllocError)?
-                .as_ptr();
-            let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
-            Ok(NonNull::slice_from_raw_parts(ptr, layout.size()))
-        })
+        let raw_ptr = unsafe { self.alloc(layout) };
+
+        if raw_ptr.is_null() {
+            return Err(AllocError);
+        }
+
+        let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
+        Ok(NonNull::slice_from_raw_parts(ptr, layout.size()))
     }
 
     unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {