m5stack-core 0.4.1

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Global heap ownership + external PSRAM integration.
//!
//! The BSP owns the global heap so a binary never spells out
//! `esp_alloc::heap_allocator!` or the per-board region sizes itself:
//! [`init_heap`] declares the esp-alloc DRAM regions for a [`HeapProfile`] and,
//! optionally, registers external PSRAM. esp-alloc's global heap holds at most
//! three regions, so each profile registers the reclaimed-ROM region, the
//! plain-DRAM region, and (when PSRAM is supplied) the external region — never a
//! fourth (a 4th `add_region` panics silently). The sizes are the HIL-proven
//! per-board values previously copied into every binary.
//!
//! ```ignore
//! use m5stack_core::mem::{self, HeapProfile};
//!
//! // After board init, before the first allocation:
//! mem::init_heap(HeapProfile::Default, Some(board.psram)); // DRAM + PSRAM
//! mem::init_heap(HeapProfile::Lvgl, None);                  // DRAM only
//! ```
//!
//! The PSRAM-specific surface below ([`init_psram_heap`], [`psram_split`], the
//! checked [`psram_box`] / [`psram_vec`], [`PsramSafe`]) needs the `psram`
//! feature; the heap regions and [`dma_buffer`] need only `heap`.
//!
//! Both boards carry SPI PSRAM (Fire27: ~4 MB, CoreS3: ~8 MB). `esp-alloc`
//! exposes a single global heap that can be backed by several regions;
//! [`init_psram_heap`] maps the external PSRAM and registers it as one such
//! region. After that, an application can allocate from it in two ways:
//!
//! 1. **Implicitly** — once registered, the global allocator may satisfy any
//!    `alloc::vec!` / `Box` / `String` from PSRAM (internal DRAM is consumed
//!    first, then it spills to PSRAM).
//! 2. **Explicitly** — pick the region per allocation. Prefer the *checked*
//!    helpers [`psram_box`] / [`psram_vec`], which reject atomic-bearing types
//!    at compile time (see [`PsramSafe`]):
//!
//! For a *private* region handed to a foreign allocator (not the global heap at
//! all) — e.g. LVGL's TLSF — use [`psram_split`] instead, which carves a private
//! slice off the base and registers only the remainder globally.
//!
//! ```ignore
//! use m5stack_core::mem;
//!
//! let psram_free = mem::init_psram_heap(peripherals.PSRAM);
//!
//! let mut big = mem::psram_vec::<u8>(512 * 1024);   // in PSRAM, atomics rejected
//! let scratch = mem::psram_box([0u32; 1024]);       // in PSRAM
//! let dma = mem::dma_buffer(4 * 1024);              // in internal DRAM, DMA-safe
//! ```
//!
//! The raw marker allocators ([`ExternalMemory`] / [`InternalMemory`]) are also
//! re-exported as an escape hatch for `allocator_api2` containers, but they do
//! **not** perform the atomic check — reach for them only when you know what
//! you are placing in PSRAM.
//!
//! ## Enforced vs. documented caveats
//!
//! - **Atomics must not live in PSRAM.** *Enforced* on the checked path:
//!   [`psram_box`] / [`psram_vec`] bound `T: PsramSafe`, so anything holding an
//!   `Atomic*` (directly or transitively) fails to compile.
//! - **DMA from PSRAM:** the original ESP32 (Fire27) cannot DMA out of PSRAM.
//!   *Guarded* by [`assert_dma_capable`] (a `debug_assert` on Fire27, a no-op on
//!   CoreS3, which can DMA from PSRAM); use [`dma_buffer`] to get an
//!   internal-DRAM buffer in the first place.
//! - **opt-level > 0:** *Enforced* at build time — enabling the `psram` feature
//!   with `opt-level = 0` fails the build (see `build.rs`). PSRAM timing
//!   calibration is unreliable unoptimized.

use allocator_api2::vec::Vec;
pub use esp_alloc::{AnyMemory, ExternalMemory, InternalMemory};
use esp_hal::peripherals::PSRAM;
use esp_hal::ram;

#[cfg(feature = "psram")]
use allocator_api2::boxed::Box;
#[cfg(feature = "psram")]
use core::mem::MaybeUninit;

/// Heap size profile — selects the HIL-proven per-board DRAM region sizes for a
/// workload. The BSP owns the sizes so every binary gets the validated values;
/// pass the matching profile to [`init_heap`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HeapProfile {
    /// Display / I2C / WiFi-STA workloads: reclaimed-ROM + plain DRAM (+ PSRAM).
    Default,
    /// LVGL UI: the reclaimed-ROM region only (LVGL object/style pool); no PSRAM.
    Lvgl,
    /// WiFi + BLE coexistence: more controller heap (+ PSRAM). Fire27 favours the
    /// reclaimed region; CoreS3 favours plain DRAM.
    Coex,
}

/// Register the global heap regions for `profile`, plus external PSRAM when
/// `psram` is `Some`, using the HIL-proven per-board sizes. Call once, right
/// after [`crate::board::init`] / `Board::split` and before any allocation.
///
/// This is the single place a binary sets up the heap — it never calls
/// `esp_alloc::heap_allocator!` itself. esp-alloc's global heap holds at most
/// three regions; each profile registers at most the reclaimed-ROM region, the
/// plain-DRAM region and the external PSRAM region — never a fourth (a 4th
/// `add_region` panics silently). Pass `None` for heap-only workloads
/// (e.g. [`HeapProfile::Lvgl`], or a board with no external RAM).
///
/// Registering PSRAM needs the `psram` feature; without it a `Some(_)` argument
/// is accepted but the external region is **not** added (the DRAM regions still
/// are).
pub fn init_heap(profile: HeapProfile, psram: Option<PSRAM<'static>>) {
    match profile {
        HeapProfile::Default => {
            esp_alloc::heap_allocator!(#[ram(reclaimed)] size: 50 * 1024);
            esp_alloc::heap_allocator!(size: 64 * 1024);
        }
        HeapProfile::Lvgl => {
            esp_alloc::heap_allocator!(#[ram(reclaimed)] size: 50 * 1024);
        }
        HeapProfile::Coex => {
            #[cfg(feature = "fire27")]
            {
                esp_alloc::heap_allocator!(#[ram(reclaimed)] size: 96 * 1024);
                esp_alloc::heap_allocator!(size: 24 * 1024);
            }
            #[cfg(feature = "cores3")]
            {
                esp_alloc::heap_allocator!(#[ram(reclaimed)] size: 50 * 1024);
                esp_alloc::heap_allocator!(size: 96 * 1024);
            }
        }
    }
    if let Some(psram) = psram {
        // `psram` feature off → the binding is consumed here without registering
        // an external region (the PSRAM controller stays uninitialised).
        let _ = &psram;
        #[cfg(feature = "psram")]
        init_psram_heap(psram);
    }
}

/// Map the board's external PSRAM and add it to the global heap as an
/// [`ExternalMemory`] region.
///
/// The size is auto-detected. Returns the amount of external (PSRAM) heap free
/// immediately after registration, in bytes.
///
/// Call once, after [`esp_hal::init`]. Usually invoked for you by [`init_heap`]
/// when you pass `Some(psram)`; call it directly only if you manage the DRAM
/// regions yourself. Calling it more than once is unsound — the PSRAM
/// controller must only be initialized a single time.
#[cfg(feature = "psram")]
pub fn init_psram_heap(psram: PSRAM<'static>) -> usize {
    esp_alloc::psram_allocator!(psram, esp_hal::psram);
    let free = esp_alloc::HEAP.free_caps(esp_alloc::MemoryCapability::External.into());
    info!("PSRAM heap registered: {} KiB external free", free / 1024);
    free
}

/// A private PSRAM region carved off the global heap, plus the external bytes
/// registered with the global heap. Returned by [`psram_split`].
#[cfg(feature = "psram")]
pub struct PsramSplit {
    /// A private, exclusive, contiguous PSRAM region for a *foreign* allocator
    /// (e.g. LVGL's built-in TLSF via `lv_mem_add_pool`). It is **not** part of
    /// the global heap, so the global allocator never hands it to `Box` / `Vec`
    /// / DMA. Its base is the PSRAM mapping base, so it is large-aligned (≥ any
    /// reasonable `ALIGN_SIZE`) — the caller needs no alignment math and no
    /// `unsafe`.
    ///
    /// `'static` is sound because esp-hal's `Psram` has no `Drop`: the mapping is
    /// a hardware side effect recorded in esp-hal's range statics and is *not*
    /// undone when the `Psram` value drops.
    pub private: &'static mut [MaybeUninit<u8>],
    /// External (PSRAM) heap free immediately after registering the remainder
    /// with the global heap, in bytes. `0` when `reserve` was `None` (all
    /// private, nothing registered).
    pub global_free: usize,
}

// `private` is a `&'static mut` to uninit memory, so `PsramSplit` cannot derive
// `Debug`; a manual impl prints the base/len/free a consumer wants when bringing
// this up on a new board (`log::info!("{:?}", split)`).
#[cfg(feature = "psram")]
impl core::fmt::Debug for PsramSplit {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("PsramSplit")
            .field("private_base", &self.private.as_ptr())
            .field("private_len", &self.private.len())
            .field("global_free", &self.global_free)
            .finish()
    }
}

/// Why [`psram_split`] could not satisfy the request.
#[cfg(feature = "psram")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PsramSplitError {
    /// PSRAM did not map — `esp_hal::psram::init_psram` failed, or the board has
    /// none. Nothing was registered with the global heap.
    NotMapped,
    /// PSRAM mapped, but smaller than the requested `reserve`. `available` is the
    /// total mapped size, so the caller can retry with a smaller pool without
    /// re-querying — the `PSRAM` peripheral was consumed by value.
    TooSmall { available: usize },
}

/// Map the board's external PSRAM, carve a **private** region off the base, and
/// register the remainder with the global heap.
///
/// The split-out counterpart of [`init_psram_heap`] (which maps *and* registers
/// the whole region globally in one step). Reach for `psram_split` when a
/// consumer needs a private, exclusive, contiguous PSRAM region to hand to a
/// *foreign* allocator — LVGL's built-in TLSF via `lv_mem_add_pool` — rather than
/// routing those allocations through the shared global allocator.
///
/// - `reserve: Some(n)` carves `n` bytes private from the base and registers the
///   remainder globally. `reserve: None` makes the whole region private (nothing
///   is added to the global heap; [`PsramSplit::global_free`] is `0`).
///   `reserve: Some(0)` is the mirror image: an empty [`PsramSplit::private`]
///   slice and *all* PSRAM registered globally — i.e. it degenerates to what
///   [`init_psram_heap`] does. Sound (a zero-length slice grants access to
///   nothing) but rarely what you want.
/// - The private region is carved **from the base**, so [`PsramSplit::private`]
///   starts at the (large-aligned) PSRAM mapping base — aligned for LVGL's TLSF
///   with no math; esp-alloc aligns the remainder's base internally.
/// - This primitive controls *placement* (the private/global split). The PSRAM
///   **hardware** mapping uses the default [`esp_hal::psram`] config, which
///   auto-detects size — the board-correct choice for both boards. A
///   `psram_split_with(config)` variant is a non-breaking addition if a consumer
///   ever needs a custom `PsramConfig` (a fixed `PsramSize`, say); no current
///   one does.
///
/// Call once, after [`esp_hal::init`], **instead of** passing `Some(psram)` to
/// [`init_heap`] or calling [`init_psram_heap`]: taking `PSRAM<'static>` by value
/// makes the once-only mapping a type-level guarantee, so the two cannot both run.
///
/// # Caveats handed back to the caller
/// - **No atomics in the private region.** The checked [`psram_box`] /
///   [`psram_vec`] cannot guard a foreign allocator, so keeping `Atomic*` out of
///   whatever is placed here is the caller's responsibility (holds for LVGL while
///   `LV_USE_OS` is `LV_OS_NONE`). See [`PsramSafe`].
/// - **DMA.** The ESP32 (Fire27) cannot DMA to/from PSRAM at all; the ESP32-S3
///   can but slowly. A foreign allocator must not place DMA'd buffers here. See
///   [`assert_dma_capable`] / [`dma_buffer`].
///
/// # Errors
/// [`PsramSplitError::NotMapped`] if PSRAM does not map; [`PsramSplitError::TooSmall`]
/// if it maps smaller than `reserve`.
#[cfg(feature = "psram")]
pub fn psram_split(
    psram: PSRAM<'static>,
    reserve: Option<usize>,
) -> Result<PsramSplit, PsramSplitError> {
    // `Psram` has no `Drop`: the mapping is recorded in esp-hal's range statics
    // and survives the value dropping (see `PsramSplit::private`), so a local is
    // fine — nothing unmaps at the end of this block.
    let psram = esp_hal::psram::Psram::new(psram, Default::default());
    let (base, total) = psram.raw_parts();

    // `Psram::new` maps only if `init_psram` succeeded; on failure the range
    // statics stay unset and `raw_parts` reports a zero-size region.
    if total == 0 {
        return Err(PsramSplitError::NotMapped);
    }

    let reserve = reserve.unwrap_or(total);
    if reserve > total {
        return Err(PsramSplitError::TooSmall { available: total });
    }

    // Carve `[base, base + reserve)` private.
    // SAFETY: `base` is the exclusively-owned, `'static` PSRAM mapping (once-only,
    // guaranteed by consuming `PSRAM<'static>` by value). This sub-range is handed
    // out privately and is never registered with the global heap, so no aliasing.
    // `MaybeUninit<u8>` has align 1, and `reserve <= total <= isize::MAX`. When
    // `reserve == 0` the slice is empty: its base pointer lies inside the region
    // that *is* registered globally below, but a zero-length slice dereferences
    // nothing, so it still aliases nothing.
    let private =
        unsafe { core::slice::from_raw_parts_mut(base as *mut MaybeUninit<u8>, reserve) };

    // Register the remainder `[base + reserve, base + total)` with the global heap.
    let global_free = if reserve < total {
        // SAFETY: disjoint from `private`, `'static`, exclusively the heap's, and
        // `total - reserve > 0` here (so esp-alloc's `size > 0` precondition holds).
        unsafe {
            esp_alloc::HEAP.add_region(esp_alloc::HeapRegion::new(
                base.add(reserve),
                total - reserve,
                esp_alloc::MemoryCapability::External.into(),
            ));
        }
        esp_alloc::HEAP.free_caps(esp_alloc::MemoryCapability::External.into())
    } else {
        0
    };

    info!(
        "PSRAM split: {} KiB private, {} KiB external free (global)",
        reserve / 1024,
        global_free / 1024
    );
    Ok(PsramSplit { private, global_free })
}

/// Marker for types safe to store in PSRAM: nothing holding an *inline* atomic.
///
/// Atomic read-modify-write instructions misbehave against PSRAM-backed
/// addresses on ESP32 / ESP32-S3, so the checked allocators [`psram_box`] /
/// [`psram_vec`] only accept `T: PsramSafe`. Like `Send` / `Sync` this is an
/// auto trait: a struct is `PsramSafe` iff every field is, so a type that
/// embeds an `Atomic*` (directly or transitively — e.g. via `Arc`, many lock
/// types) is rejected at compile time.
///
/// A *pointer or reference* to an atomic living elsewhere is fine — the atomic
/// itself is not in PSRAM — so `&T`, `&mut T`, `*const T` and `*mut T` are
/// always `PsramSafe`.
///
/// # Safety
/// Only implement (or negative-impl) this to reflect the atomic-in-PSRAM
/// hazard; the checked allocators rely on it to keep atomics out of PSRAM.
#[cfg(feature = "psram")]
pub unsafe auto trait PsramSafe {}

#[cfg(feature = "psram")]
mod psram_safe_impls {
    use super::PsramSafe;
    use core::sync::atomic::{
        AtomicBool, AtomicI8, AtomicI16, AtomicI32, AtomicIsize, AtomicPtr, AtomicU8, AtomicU16,
        AtomicU32, AtomicUsize,
    };

    impl !PsramSafe for AtomicBool {}
    impl !PsramSafe for AtomicI8 {}
    impl !PsramSafe for AtomicU8 {}
    impl !PsramSafe for AtomicI16 {}
    impl !PsramSafe for AtomicU16 {}
    impl !PsramSafe for AtomicI32 {}
    impl !PsramSafe for AtomicU32 {}
    impl !PsramSafe for AtomicIsize {}
    impl !PsramSafe for AtomicUsize {}
    impl<T> !PsramSafe for AtomicPtr<T> {}

    // A pointer/reference to an atomic is fine — the atomic lives elsewhere.
    // (Mirrors `unsafe impl<T: ?Sized> Send for &T` in std.)
    unsafe impl<T: ?Sized> PsramSafe for &T {}
    unsafe impl<T: ?Sized> PsramSafe for &mut T {}
    unsafe impl<T: ?Sized> PsramSafe for *const T {}
    unsafe impl<T: ?Sized> PsramSafe for *mut T {}
}

/// Allocate `value` in external PSRAM. Atomic-bearing `T` is rejected at
/// compile time via [`PsramSafe`].
#[cfg(feature = "psram")]
pub fn psram_box<T: PsramSafe>(value: T) -> Box<T, ExternalMemory> {
    Box::new_in(value, ExternalMemory)
}

/// A `Vec<T>` with room for `capacity` elements reserved in external PSRAM.
/// Atomic-bearing `T` is rejected at compile time via [`PsramSafe`].
#[cfg(feature = "psram")]
pub fn psram_vec<T: PsramSafe>(capacity: usize) -> Vec<T, ExternalMemory> {
    Vec::with_capacity_in(capacity, ExternalMemory)
}

/// Free bytes in the global heap's **internal** DRAM regions (reclaimed-ROM +
/// plain-DRAM), right now. The tight resource on both boards; use it to measure
/// headroom (e.g. before/after moving a subsystem's heap to PSRAM).
pub fn internal_free() -> usize {
    esp_alloc::HEAP.free_caps(esp_alloc::MemoryCapability::Internal.into())
}

/// Free bytes in the global heap's **external** (PSRAM) region, right now — `0`
/// unless PSRAM was registered globally (via [`init_heap`] with `Some(psram)` /
/// [`init_psram_heap`]; a private [`psram_split`] region is *not* counted here).
pub fn external_free() -> usize {
    esp_alloc::HEAP.free_caps(esp_alloc::MemoryCapability::External.into())
}

/// A zeroed byte buffer in internal DRAM, suitable as a DMA buffer.
///
/// Convenience for the common "I need a DMA-capable scratch buffer" case so the
/// allocator does not have to be spelled out. The result is DMA-reachable on
/// both chips; pair it with [`assert_dma_capable`] if a buffer's origin is ever
/// in doubt.
pub fn dma_buffer(len: usize) -> Vec<u8, InternalMemory> {
    let mut v = Vec::with_capacity_in(len, InternalMemory);
    v.resize(len, 0);
    v
}

/// Debug-assert that `buf` is DMA-reachable on this chip.
///
/// On the ESP32 (Fire27) the DMA engine cannot reach the PSRAM-mapped data
/// window, so a PSRAM-backed buffer handed to SPI/I2S DMA silently corrupts.
/// This catches that on first use under `debug_assertions`. It is a no-op
/// (compiled away) on the ESP32-S3, which *can* DMA from PSRAM.
#[cfg(feature = "fire27")]
#[inline]
pub fn assert_dma_capable(buf: &[u8]) {
    // ESP32 external RAM (PSRAM) is cache-mapped into this data window; internal
    // DRAM lives above it. Constant per the ESP32 TRM external-memory map.
    const PSRAM_DATA_WINDOW: core::ops::Range<usize> = 0x3F80_0000..0x3FC0_0000;
    let p = buf.as_ptr() as usize;
    debug_assert!(
        !PSRAM_DATA_WINDOW.contains(&p),
        "DMA buffer at {p:#x} lives in PSRAM; the ESP32 cannot DMA to/from PSRAM \
         — allocate it in InternalMemory (see mem::dma_buffer)"
    );
}

/// No-op on every target except the ESP32 (Fire27); see the Fire27 variant.
#[cfg(not(feature = "fire27"))]
#[inline]
pub fn assert_dma_capable(_buf: &[u8]) {}