From 506b0e8384e7494cc9210706f77de0f45557811a Mon Sep 17 00:00:00 2001
From: Markus Ebner <info@ebner-markus.de>
Date: Sat, 2 Mar 2024 21:02:08 +0100
Subject: [PATCH] Add support for nightly allocator_api-feature

---
 Cargo.toml |  3 +++
 src/lib.rs | 65 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/Cargo.toml b/Cargo.toml
index 6319522..694c4e2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -22,6 +22,9 @@ categories = [
     "no-std",
 ]
 
+[features]
+nightly = []
+
 [package.metadata.docs.rs]
 default-target = "riscv32imc-unknown-none-elf"
 
diff --git a/src/lib.rs b/src/lib.rs
index 4ff3667..24956d8 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,8 +3,49 @@
 //!
 //! **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();
+//!
+//! 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);
+//!     }
+//! }
+//! ```
+//!
+//! # 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, which allows
+//! you to specify an allocator for single allocations.
+//!
+//! **NOTE:** To use this, you have to enable the create's `nightly` feature flag.
+//!
+//! Create and initialize an allocator to use in single allocations:
+//! ```rust
+//! static PSRAM_ALLOCATOR: esp_alloc::EspHeap = esp_alloc::EspHeap::empty();
+//!
+//! fn init_psram_heap() {
+//!     unsafe {
+//!         PSRAM_ALLOCATOR.init(psram::psram_vaddr_start() as *mut u8, psram::PSRAM_BYTES);
+//!     }
+//! }
+//! ```
+//!
+//! And then use it in an allocation:
+//! ```rust
+//! let large_buffer: Vec<u8, _> = Vec::with_capacity_in(1048576, &PSRAM_ALLOCATOR);
+//! ```
 
 #![no_std]
+#![cfg_attr(feature = "nightly", feature(allocator_api))]
 
 pub mod macros;
 
@@ -14,6 +55,9 @@ use core::{
     ptr::{self, NonNull},
 };
 
+#[ cfg( feature = "nightly" ) ]
+use core::alloc::{AllocError, Allocator};
+
 use critical_section::Mutex;
 use linked_list_allocator::Heap;
 
@@ -92,3 +136,24 @@ unsafe impl GlobalAlloc for EspHeap {
         });
     }
 }
+
+
+#[ 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()))
+        })
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        self.dealloc(ptr.as_ptr(), layout);
+    }
+}