rublon-ssh/PAM/ssh/include/rublon/memory.hpp

#pragma once

#include <memory_resource>

namespace rublon {
namespace memory {
    struct default_memory_resource {
        static inline std::pmr::memory_resource * _mr = std::pmr::get_default_resource();
    };

    inline void set_default_resource(std::pmr::memory_resource * memory_resource) {
        default_memory_resource{}._mr = memory_resource;
    }

    inline std::pmr::memory_resource * default_resource() {
        return default_memory_resource{}._mr;
    }

    template < std::size_t N >
    class MonotonicStackResource : public std::pmr::monotonic_buffer_resource {
        char _buffer[N];

      public:
        MonotonicStackResource() : std::pmr::monotonic_buffer_resource{_buffer, N, std::pmr::null_memory_resource()} {}
    };

    template < std::size_t N >
    class UnsynchronizedStackResource : public std::pmr::unsynchronized_pool_resource {
        MonotonicStackResource< N > _upstream;

      public:
        UnsynchronizedStackResource() : std::pmr::unsynchronized_pool_resource{&_upstream} {}
    };

    class MonotonicHeapResourceBase {
      public:
        std::pmr::memory_resource * _upstream{};
        std::size_t _size{};
        void * _buffer{nullptr};

        MonotonicHeapResourceBase(std::size_t size) : _upstream{default_resource()}, _size{size}, _buffer{_upstream->allocate(size)} {}

        ~MonotonicHeapResourceBase() {
            if(_buffer)
                _upstream->deallocate(_buffer, _size);
        }
    };

    template < std::size_t N >
    class MonotonicHeapResource : MonotonicHeapResourceBase, public std::pmr::monotonic_buffer_resource {
      public:
        MonotonicHeapResource()
         : MonotonicHeapResourceBase{N}, std::pmr::monotonic_buffer_resource{this->_buffer, this->_size, default_resource()} {}
    };

    template < std::size_t N >
    class StrictMonotonicHeapResource : MonotonicHeapResourceBase, public std::pmr::monotonic_buffer_resource {
      public:
        StrictMonotonicHeapResource()
         : MonotonicHeapResourceBase{N},
           std::pmr::monotonic_buffer_resource{this->_buffer, this->_size, std::pmr::null_memory_resource()} {}
    };
    
    using StrictMonotonic_512_HeapResource = StrictMonotonicHeapResource< 512 >;
    using StrictMonotonic_1k_HeapResource = StrictMonotonicHeapResource< 1 * 1024 >;
    using StrictMonotonic_2k_HeapResource = StrictMonotonicHeapResource< 2 * 1024 >;
    using StrictMonotonic_4k_HeapResource = StrictMonotonicHeapResource< 4 * 1024 >;
    using StrictMonotonic_8k_HeapResource = StrictMonotonicHeapResource< 8 * 1024 >;

    using Monotonic_1k_HeapResource = MonotonicHeapResource< 1 * 1024 >;
    using Monotonic_2k_HeapResource = MonotonicHeapResource< 2 * 1024 >;
    using Monotonic_4k_HeapResource = MonotonicHeapResource< 4 * 1024 >;
    using Monotonic_8k_HeapResource = MonotonicHeapResource< 8 * 1024 >;
} // namespace memory

class RublonMemory {
  public:
    std::byte sharedMemory[32 * 1024] = {};
    std::pmr::monotonic_buffer_resource mr{sharedMemory, std::size(sharedMemory)};
    std::pmr::unsynchronized_pool_resource rublonPoolResource{&mr};

    RublonMemory() {    
        std::pmr::set_default_resource(&rublonPoolResource);
    }
};
} // namespace rublon