AllocatorPSRAM.h

#pragma once
#include "assert.h"
 
#ifdef ESP32
#include <esp_heap_caps.h>
#endif
 
namespace tinyrobotics {
 
/**
 * @class AllocatorPSRAM
 * @ingroup utils
 * @brief Custom allocator that uses ESP32's PSRAM for memory allocation
 * @tparam T Type of elements to allocate
 *
 * This allocator uses ESP32's heap_caps_malloc with MALLOC_CAP_SPIRAM flag
 * to ensure all memory is allocated in PSRAM instead of regular RAM. If PSRAM
 * allocation fails, it falls back to regular RAM.
 */
template <typename T>
class AllocatorPSRAM {
 public:
  using value_type = T;
  using pointer = T*;
  using const_pointer = const T*;
  using reference = T&;
  using const_reference = const T&;
  using size_type = std::size_t;
  using difference_type = std::ptrdiff_t;
 
  /**
   * @brief Default constructor
   */
  AllocatorPSRAM() noexcept {}
 
  /**
   * @brief Copy constructor from another allocator type
   * @tparam U Type of the other allocator
   * @param other The other allocator
   */
  template <typename U>
  AllocatorPSRAM(const AllocatorPSRAM<U>&) noexcept {}
 
  /**
   * @brief Allocate memory from PSRAM
   * @param n Number of elements to allocate
   * @return Pointer to allocated memory
   * @throws std::bad_alloc If allocation fails or size is too large
   */
  pointer allocate(size_type n) {
    // if (n > std::numeric_limits<size_type>::max() / sizeof(T))
    //     throw std::bad_alloc();
    // in Arduino excepitons are disabled!
    assert(n <= std::numeric_limits<size_type>::max() / sizeof(T));
 
    pointer p = nullptr;
#ifdef ESP32
    p = static_cast<pointer>(
        heap_caps_malloc(n * sizeof(T), MALLOC_CAP_SPIRAM));
#endif
#if defined(ARDUINO_ARCH_RP2040) && defined(RP2350_PSRAM_CS)
    p = pmalloc(n * sizeof(T));
#endif
 
    if (p == nullptr) {
      p = static_cast<pointer>(malloc(n * sizeof(T)));
    }
 
    // if (!p) throw std::bad_alloc();
 
    // in Arduino excepitons are disabled!
    assert(p);
 
    return p;
  }
 
  /**
   * @brief Deallocate memory
   * @param p Pointer to memory to deallocate
   * @param size Size of allocation (unused)
   */
  void deallocate(pointer p, size_type) noexcept {
#ifdef ESP32
    heap_caps_free(p);
#else
    free(p);
#endif
  }
 
  // Required for STL allocator compatibility
  template <typename U>
  bool operator==(const AllocatorPSRAM<U>&) const noexcept {
    return true;
  }
  template <typename U>
  bool operator!=(const AllocatorPSRAM<U>&) const noexcept {
    return false;
  }
  /**
   * @brief Rebind allocator to another type
   * @tparam U Type to rebind the allocator to
   */
  template <typename U>
  struct rebind {
    using other = AllocatorPSRAM<U>;
  };
};
 
}  // namespace tinyrobotics