esp_alloc/
lib.rs

1//! A `no_std` heap allocator for RISC-V and Xtensa processors from
2//! Espressif. Supports all currently available ESP32 devices.
3//!
4//! **NOTE:** using this as your global allocator requires using Rust 1.68 or
5//! greater, or the `nightly` release channel.
6//!
7//! # Using this as your Global Allocator
8//!
9//! ```rust
10//! use esp_alloc as _;
11//!
12//! fn init_heap() {
13//!     const HEAP_SIZE: usize = 32 * 1024;
14//!     static mut HEAP: MaybeUninit<[u8; HEAP_SIZE]> = MaybeUninit::uninit();
15//!
16//!     unsafe {
17//!         esp_alloc::HEAP.add_region(esp_alloc::HeapRegion::new(
18//!             HEAP.as_mut_ptr() as *mut u8,
19//!             HEAP_SIZE,
20//!             esp_alloc::MemoryCapability::Internal.into(),
21//!         ));
22//!     }
23//! }
24//! ```
25//!
26//! # Using this with the nightly `allocator_api`-feature
27//! Sometimes you want to have more control over allocations.
28//!
29//! For that, it's convenient to use the nightly `allocator_api`-feature,
30//! which allows you to specify an allocator for single allocations.
31//!
32//! **NOTE:** To use this, you have to enable the crate's `nightly` feature
33//! flag.
34//!
35//! Create and initialize an allocator to use in single allocations:
36//! ```rust
37//! static PSRAM_ALLOCATOR: esp_alloc::EspHeap = esp_alloc::EspHeap::empty();
38//!
39//! fn init_psram_heap() {
40//!     unsafe {
41//!         PSRAM_ALLOCATOR.add_region(esp_alloc::HeapRegion::new(
42//!             psram::psram_vaddr_start() as *mut u8,
43//!             psram::PSRAM_BYTES,
44//!             esp_alloc::MemoryCapability::Internal.into(),
45//!         ));
46//!     }
47//! }
48//! ```
49//!
50//! And then use it in an allocation:
51//! ```rust
52//! let large_buffer: Vec<u8, _> = Vec::with_capacity_in(1048576, &PSRAM_ALLOCATOR);
53//! ```
54//!
55//! You can also get stats about the heap usage at anytime with:
56//! ```rust
57//! let stats: HeapStats = esp_alloc::HEAP.stats();
58//! // HeapStats implements the Display and defmt::Format traits, so you can pretty-print the heap stats.
59//! println!("{}", stats);
60//! ```
61//!
62//! ```txt
63//! HEAP INFO
64//! Size: 131068
65//! Current usage: 46148
66//! Max usage: 46148
67//! Total freed: 0
68//! Total allocated: 46148
69//! Memory Layout:
70//! Internal | ████████████░░░░░░░░░░░░░░░░░░░░░░░ | Used: 35% (Used 46148 of 131068, free: 84920)
71//! Unused   | ░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░ |
72//! Unused   | ░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░ |
73//! ```
74//! ## Feature Flags
75#![doc = document_features::document_features!()]
76#![no_std]
77#![cfg_attr(feature = "nightly", feature(allocator_api))]
78#![doc(html_logo_url = "https://avatars.githubusercontent.com/u/46717278")]
79
80mod macros;
81
82#[cfg(feature = "nightly")]
83use core::alloc::{AllocError, Allocator};
84use core::{
85    alloc::{GlobalAlloc, Layout},
86    cell::RefCell,
87    fmt::Display,
88    ptr::{self, NonNull},
89};
90
91use critical_section::Mutex;
92use enumset::{EnumSet, EnumSetType};
93use linked_list_allocator::Heap;
94
95/// The global allocator instance
96#[global_allocator]
97pub static HEAP: EspHeap = EspHeap::empty();
98
99const NON_REGION: Option<HeapRegion> = None;
100
101const BAR_WIDTH: usize = 35;
102
103fn write_bar(f: &mut core::fmt::Formatter<'_>, usage_percent: usize) -> core::fmt::Result {
104    let used_blocks = BAR_WIDTH * usage_percent / 100;
105    (0..used_blocks).try_for_each(|_| write!(f, "█"))?;
106    (used_blocks..BAR_WIDTH).try_for_each(|_| write!(f, "░"))
107}
108
109#[cfg(feature = "defmt")]
110fn write_bar_defmt(fmt: defmt::Formatter, usage_percent: usize) {
111    let used_blocks = BAR_WIDTH * usage_percent / 100;
112    (0..used_blocks).for_each(|_| defmt::write!(fmt, "█"));
113    (used_blocks..BAR_WIDTH).for_each(|_| defmt::write!(fmt, "░"));
114}
115
116#[derive(EnumSetType, Debug)]
117/// Describes the properties of a memory region
118pub enum MemoryCapability {
119    /// Memory must be internal; specifically it should not disappear when
120    /// flash/spiram cache is switched off
121    Internal,
122    /// Memory must be in SPI RAM
123    External,
124}
125
126/// Stats for a heap region
127#[derive(Debug)]
128pub struct RegionStats {
129    /// Total usable size of the heap region in bytes.
130    size: usize,
131
132    /// Currently used size of the heap region in bytes.
133    used: usize,
134
135    /// Free size of the heap region in bytes.
136    free: usize,
137
138    /// Capabilities of the memory region.
139    capabilities: EnumSet<MemoryCapability>,
140}
141
142impl Display for RegionStats {
143    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
144        let usage_percent = self.used * 100 / self.size;
145
146        // Display Memory type
147        if self.capabilities.contains(MemoryCapability::Internal) {
148            write!(f, "Internal")?;
149        } else if self.capabilities.contains(MemoryCapability::External) {
150            write!(f, "External")?;
151        } else {
152            write!(f, "Unknown")?;
153        }
154
155        write!(f, " | ")?;
156
157        write_bar(f, usage_percent)?;
158
159        write!(
160            f,
161            " | Used: {}% (Used {} of {}, free: {})",
162            usage_percent, self.used, self.size, self.free
163        )
164    }
165}
166
167#[cfg(feature = "defmt")]
168impl defmt::Format for RegionStats {
169    fn format(&self, fmt: defmt::Formatter<'_>) {
170        let usage_percent = self.used * 100 / self.size;
171
172        if self.capabilities.contains(MemoryCapability::Internal) {
173            defmt::write!(fmt, "Internal");
174        } else if self.capabilities.contains(MemoryCapability::External) {
175            defmt::write!(fmt, "External");
176        } else {
177            defmt::write!(fmt, "Unknown");
178        }
179
180        defmt::write!(fmt, " | ");
181
182        write_bar_defmt(fmt, usage_percent);
183
184        defmt::write!(
185            fmt,
186            " | Used: {}% (Used {} of {}, free: {})",
187            usage_percent,
188            self.used,
189            self.size,
190            self.free
191        );
192    }
193}
194
195/// A memory region to be used as heap memory
196pub struct HeapRegion {
197    heap: Heap,
198    capabilities: EnumSet<MemoryCapability>,
199}
200
201impl HeapRegion {
202    /// Create a new [HeapRegion] with the given capabilities
203    ///
204    /// # Safety
205    ///
206    /// - The supplied memory region must be available for the entire program
207    ///   (`'static`).
208    /// - The supplied memory region must be exclusively available to the heap
209    ///   only, no aliasing.
210    /// - `size > 0`.
211    pub unsafe fn new(
212        heap_bottom: *mut u8,
213        size: usize,
214        capabilities: EnumSet<MemoryCapability>,
215    ) -> Self {
216        let mut heap = Heap::empty();
217        heap.init(heap_bottom, size);
218
219        Self { heap, capabilities }
220    }
221
222    /// Return stats for the current memory region
223    pub fn stats(&self) -> RegionStats {
224        RegionStats {
225            size: self.heap.size(),
226            used: self.heap.used(),
227            free: self.heap.free(),
228            capabilities: self.capabilities,
229        }
230    }
231}
232
233/// Stats for a heap allocator
234///
235/// Enable the "internal-heap-stats" feature if you want collect additional heap
236/// informations at the cost of extra cpu time during every alloc/dealloc.
237#[derive(Debug)]
238pub struct HeapStats {
239    /// Granular stats for all the configured memory regions.
240    region_stats: [Option<RegionStats>; 3],
241
242    /// Total size of all combined heap regions in bytes.
243    size: usize,
244
245    /// Current usage of the heap across all configured regions in bytes.
246    current_usage: usize,
247
248    /// Estimation of the max used heap in bytes.
249    #[cfg(feature = "internal-heap-stats")]
250    max_usage: usize,
251
252    /// Estimation of the total allocated bytes since initialization.
253    #[cfg(feature = "internal-heap-stats")]
254    total_allocated: usize,
255
256    /// Estimation of the total freed bytes since initialization.
257    #[cfg(feature = "internal-heap-stats")]
258    total_freed: usize,
259}
260
261impl Display for HeapStats {
262    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
263        writeln!(f, "HEAP INFO")?;
264        writeln!(f, "Size: {}", self.size)?;
265        writeln!(f, "Current usage: {}", self.current_usage)?;
266        #[cfg(feature = "internal-heap-stats")]
267        {
268            writeln!(f, "Max usage: {}", self.max_usage)?;
269            writeln!(f, "Total freed: {}", self.total_freed)?;
270            writeln!(f, "Total allocated: {}", self.total_allocated)?;
271        }
272        writeln!(f, "Memory Layout: ")?;
273        for region in self.region_stats.iter() {
274            if let Some(region) = region.as_ref() {
275                region.fmt(f)?;
276                writeln!(f)?;
277            } else {
278                // Display unused memory regions
279                write!(f, "Unused   | ")?;
280                write_bar(f, 0)?;
281                writeln!(f, " |")?;
282            }
283        }
284        Ok(())
285    }
286}
287
288#[cfg(feature = "defmt")]
289impl defmt::Format for HeapStats {
290    fn format(&self, fmt: defmt::Formatter<'_>) {
291        defmt::write!(fmt, "HEAP INFO\n");
292        defmt::write!(fmt, "Size: {}\n", self.size);
293        defmt::write!(fmt, "Current usage: {}\n", self.current_usage);
294        #[cfg(feature = "internal-heap-stats")]
295        {
296            defmt::write!(fmt, "Max usage: {}\n", self.max_usage);
297            defmt::write!(fmt, "Total freed: {}\n", self.total_freed);
298            defmt::write!(fmt, "Total allocated: {}\n", self.total_allocated);
299        }
300        defmt::write!(fmt, "Memory Layout:\n");
301        for region in self.region_stats.iter() {
302            if let Some(region) = region.as_ref() {
303                defmt::write!(fmt, "{}\n", region);
304            } else {
305                defmt::write!(fmt, "Unused   | ");
306                write_bar_defmt(fmt, 0);
307                defmt::write!(fmt, " |\n");
308            }
309        }
310    }
311}
312
313/// Internal stats to keep track across multiple regions.
314#[cfg(feature = "internal-heap-stats")]
315struct InternalHeapStats {
316    max_usage: usize,
317    total_allocated: usize,
318    total_freed: usize,
319}
320
321/// A memory allocator
322///
323/// In addition to what Rust's memory allocator can do it allows to allocate
324/// memory in regions satisfying specific needs.
325pub struct EspHeap {
326    heap: Mutex<RefCell<[Option<HeapRegion>; 3]>>,
327    #[cfg(feature = "internal-heap-stats")]
328    internal_heap_stats: Mutex<RefCell<InternalHeapStats>>,
329}
330
331impl EspHeap {
332    /// Crate a new UNINITIALIZED heap allocator
333    pub const fn empty() -> Self {
334        EspHeap {
335            heap: Mutex::new(RefCell::new([NON_REGION; 3])),
336            #[cfg(feature = "internal-heap-stats")]
337            internal_heap_stats: Mutex::new(RefCell::new(InternalHeapStats {
338                max_usage: 0,
339                total_allocated: 0,
340                total_freed: 0,
341            })),
342        }
343    }
344
345    /// Add a memory region to the heap
346    ///
347    /// `heap_bottom` is a pointer to the location of the bottom of the heap.
348    ///
349    /// `size` is the size of the heap in bytes.
350    ///
351    /// You can add up to three regions per allocator.
352    ///
353    /// Note that:
354    ///
355    /// - Memory is allocated from the first suitable memory region first
356    ///
357    /// - The heap grows "upwards", towards larger addresses. Thus `end_addr`
358    ///   must be larger than `start_addr`
359    ///
360    /// - The size of the heap is `(end_addr as usize) - (start_addr as usize)`.
361    ///   The allocator won't use the byte at `end_addr`.
362    ///
363    /// # Safety
364    ///
365    /// - The supplied memory region must be available for the entire program (a
366    ///   `'static` lifetime).
367    /// - The supplied memory region must be exclusively available to the heap
368    ///   only, no aliasing.
369    /// - `size > 0`.
370    pub unsafe fn add_region(&self, region: HeapRegion) {
371        critical_section::with(|cs| {
372            let mut regions = self.heap.borrow_ref_mut(cs);
373            let free = regions
374                .iter()
375                .enumerate()
376                .find(|v| v.1.is_none())
377                .map(|v| v.0);
378
379            if let Some(free) = free {
380                regions[free] = Some(region);
381            } else {
382                panic!(
383                    "Exceeded the maximum of {} heap memory regions",
384                    regions.len()
385                );
386            }
387        });
388    }
389
390    /// Returns an estimate of the amount of bytes in use in all memory regions.
391    pub fn used(&self) -> usize {
392        critical_section::with(|cs| {
393            let regions = self.heap.borrow_ref(cs);
394            let mut used = 0;
395            for region in regions.iter() {
396                if let Some(region) = region.as_ref() {
397                    used += region.heap.used();
398                }
399            }
400            used
401        })
402    }
403
404    /// Return usage stats for the [Heap].
405    ///
406    /// Note:
407    /// [HeapStats] directly implements [Display], so this function can be
408    /// called from within `println!()` to pretty-print the usage of the
409    /// heap.
410    pub fn stats(&self) -> HeapStats {
411        const EMPTY_REGION_STAT: Option<RegionStats> = None;
412        let mut region_stats: [Option<RegionStats>; 3] = [EMPTY_REGION_STAT; 3];
413
414        critical_section::with(|cs| {
415            let mut used = 0;
416            let mut free = 0;
417            let regions = self.heap.borrow_ref(cs);
418            for (id, region) in regions.iter().enumerate() {
419                if let Some(region) = region.as_ref() {
420                    let stats = region.stats();
421                    free += stats.free;
422                    used += stats.used;
423                    region_stats[id] = Some(region.stats());
424                }
425            }
426
427            cfg_if::cfg_if! {
428                if #[cfg(feature = "internal-heap-stats")] {
429                    let internal_heap_stats = self.internal_heap_stats.borrow_ref(cs);
430                    HeapStats {
431                        region_stats,
432                        size: free + used,
433                        current_usage: used,
434                        max_usage: internal_heap_stats.max_usage,
435                        total_allocated: internal_heap_stats.total_allocated,
436                        total_freed: internal_heap_stats.total_freed,
437                    }
438                } else {
439                    HeapStats {
440                        region_stats,
441                        size: free + used,
442                        current_usage: used,
443                    }
444                }
445            }
446        })
447    }
448
449    /// Returns an estimate of the amount of bytes available.
450    pub fn free(&self) -> usize {
451        self.free_caps(EnumSet::empty())
452    }
453
454    /// The free heap satisfying the given requirements
455    pub fn free_caps(&self, capabilities: EnumSet<MemoryCapability>) -> usize {
456        critical_section::with(|cs| {
457            let regions = self.heap.borrow_ref(cs);
458            let mut free = 0;
459            for region in regions.iter().filter(|region| {
460                if region.is_some() {
461                    region
462                        .as_ref()
463                        .unwrap()
464                        .capabilities
465                        .is_superset(capabilities)
466                } else {
467                    false
468                }
469            }) {
470                if let Some(region) = region.as_ref() {
471                    free += region.heap.free();
472                }
473            }
474            free
475        })
476    }
477
478    /// Allocate memory in a region satisfying the given requirements.
479    ///
480    /// # Safety
481    ///
482    /// This function is unsafe because undefined behavior can result
483    /// if the caller does not ensure that `layout` has non-zero size.
484    ///
485    /// The allocated block of memory may or may not be initialized.
486    pub unsafe fn alloc_caps(
487        &self,
488        capabilities: EnumSet<MemoryCapability>,
489        layout: Layout,
490    ) -> *mut u8 {
491        critical_section::with(|cs| {
492            #[cfg(feature = "internal-heap-stats")]
493            let before = self.used();
494            let mut regions = self.heap.borrow_ref_mut(cs);
495            let mut iter = (*regions).iter_mut().filter(|region| {
496                if region.is_some() {
497                    region
498                        .as_ref()
499                        .unwrap()
500                        .capabilities
501                        .is_superset(capabilities)
502                } else {
503                    false
504                }
505            });
506
507            let res = loop {
508                if let Some(Some(region)) = iter.next() {
509                    let res = region.heap.allocate_first_fit(layout);
510                    if let Ok(res) = res {
511                        break Some(res);
512                    }
513                } else {
514                    break None;
515                }
516            };
517
518            res.map_or(ptr::null_mut(), |allocation| {
519                #[cfg(feature = "internal-heap-stats")]
520                {
521                    let mut internal_heap_stats = self.internal_heap_stats.borrow_ref_mut(cs);
522                    drop(regions);
523                    // We need to call used because [linked_list_allocator::Heap] does internal size
524                    // alignment so we cannot use the size provided by the layout.
525                    let used = self.used();
526
527                    internal_heap_stats.total_allocated += used - before;
528                    internal_heap_stats.max_usage =
529                        core::cmp::max(internal_heap_stats.max_usage, used);
530                }
531
532                allocation.as_ptr()
533            })
534        })
535    }
536}
537
538unsafe impl GlobalAlloc for EspHeap {
539    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
540        self.alloc_caps(EnumSet::empty(), layout)
541    }
542
543    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
544        if ptr.is_null() {
545            return;
546        }
547
548        critical_section::with(|cs| {
549            #[cfg(feature = "internal-heap-stats")]
550            let before = self.used();
551            let mut regions = self.heap.borrow_ref_mut(cs);
552            let mut iter = (*regions).iter_mut();
553
554            while let Some(Some(region)) = iter.next() {
555                if region.heap.bottom() <= ptr && region.heap.top() >= ptr {
556                    region.heap.deallocate(NonNull::new_unchecked(ptr), layout);
557                }
558            }
559
560            #[cfg(feature = "internal-heap-stats")]
561            {
562                let mut internal_heap_stats = self.internal_heap_stats.borrow_ref_mut(cs);
563                drop(regions);
564                // We need to call `used()` because [linked_list_allocator::Heap] does internal
565                // size alignment so we cannot use the size provided by the
566                // layout.
567                internal_heap_stats.total_freed += before - self.used();
568            }
569        })
570    }
571}
572
573#[cfg(feature = "nightly")]
574unsafe impl Allocator for EspHeap {
575    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
576        let raw_ptr = unsafe { self.alloc(layout) };
577
578        if raw_ptr.is_null() {
579            return Err(AllocError);
580        }
581
582        let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
583        Ok(NonNull::slice_from_raw_parts(ptr, layout.size()))
584    }
585
586    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
587        self.dealloc(ptr.as_ptr(), layout);
588    }
589}