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

#pragma once

#include "rapidjson/encodings.h"
#include <rapidjson/document.h>
#include <rapidjson/istreamwrapper.h>
#include <rapidjson/pointer.h>
#include <rapidjson/prettywriter.h>
#include <rapidjson/writer.h>
#include <rublon/utils.hpp>

#include <cstring>
namespace rublon {

struct RapidJSONPMRAlloc {
    std::pmr::memory_resource * upstream = std::pmr::get_default_resource();

    static constexpr bool kNeedFree    = true;
    static constexpr auto objectOffset = alignof(std::max_align_t);
    static constexpr auto memPadding   = objectOffset * 2;

    RapidJSONPMRAlloc(std::pmr::memory_resource * mr = std::pmr::get_default_resource()) : upstream{mr} {}
    void * Malloc(size_t size) {
        if(size != 0) {
            const auto allocated_size = size + memPadding;
            std::byte * newPtr        = static_cast< std::byte * >(upstream->allocate(allocated_size));
            auto * ptrToReturn        = newPtr + memPadding;
            // placement new a pointer to ourselves at the first memory location
            new(newPtr)(RapidJSONPMRAlloc *)(this);
            // placement new the size in the second location
            new(newPtr + objectOffset)(std::size_t)(size);
            return ptrToReturn;
        } else {
            return nullptr;
        }
    }

    void freePtr(void * origPtr, size_t originalSize) {
        if(origPtr == nullptr) {
            return;
        }
        upstream->deallocate(static_cast< std::byte * >(origPtr) - memPadding, originalSize + memPadding);
    }

    void * Realloc(void * origPtr, size_t originalSize, size_t newSize) {
        if(newSize == 0) {
            freePtr(origPtr, originalSize);
            return nullptr;
        }

        if(newSize <= originalSize) {
            return origPtr;
        }

        void * newPtr = Malloc(newSize);
        if(originalSize) {
            std::memcpy(newPtr, origPtr, originalSize);
            freePtr(origPtr, originalSize);
        }
        return newPtr;
    }

    // and Free needs to be static, which causes this whole thing
    // to fall apart. This means that we have to keep our own list of allocated memory
    // with our own pointers back to ourselves and our own list of sizes
    // so we can push all of this back to the upstream allocator
    static void Free(void * ptr) {
        if(ptr == nullptr) {
            return;
        }

        std::byte * startOfData = static_cast< std::byte * >(ptr) - memPadding;

        auto * ptrToAllocator  = *reinterpret_cast< RapidJSONPMRAlloc ** >(startOfData);
        auto origAllocatedSize = *reinterpret_cast< std::size_t * >(startOfData + objectOffset);

        ptrToAllocator->freePtr(ptr, origAllocatedSize);
    }
};

template < std::size_t N >
struct RapidJSONPMRStackAlloc : public RapidJSONPMRAlloc {
  private:
    char _buffer[N]{};
    std::pmr::monotonic_buffer_resource mr{_buffer, N};

  public:
    RapidJSONPMRStackAlloc() : RapidJSONPMRAlloc(&mr) {}
    RapidJSONPMRStackAlloc(const RapidJSONPMRStackAlloc &)             = delete;
    RapidJSONPMRStackAlloc(RapidJSONPMRStackAlloc &&)                  = delete;
    RapidJSONPMRStackAlloc & operator=(const RapidJSONPMRStackAlloc &) = delete;
    RapidJSONPMRStackAlloc & operator=(RapidJSONPMRStackAlloc &&)      = delete;
};

using Document     = rapidjson::GenericDocument< rapidjson::UTF8<>, RapidJSONPMRAlloc >;
using Value        = rapidjson::GenericValue< rapidjson::UTF8<>, RapidJSONPMRAlloc >;
using StringBuffer = rapidjson::GenericStringBuffer< rapidjson::UTF8<>, RapidJSONPMRAlloc >;
using Writer       = rapidjson::Writer< StringBuffer, rapidjson::UTF8<>, rapidjson::UTF8<>, RapidJSONPMRAlloc >;

using JSONPointer = rapidjson::GenericPointer< Value, RapidJSONPMRAlloc >;

template < typename Str_t >
struct StringWriter {
    Str_t & _str;
    using Ch = char;
    StringWriter(Str_t & t) : _str{t} {}
    void Put(Ch ch) {
        _str += ch;
    }
    void Flush() {}
};
struct FileWriter {
    std::unique_ptr< FILE, int (*)(FILE *) > _fp{nullptr, nullptr};
    using Ch = char;
    FileWriter(std::string_view filename) {
        _fp = std::unique_ptr< FILE, int (*)(FILE *) >(fopen(filename.data(), "w"), fclose);
    }
    void Put(Ch ch) {
        fputc(ch, _fp.get());
    }
    void Flush() {}
};

template < typename T >
static void stringifyTo(const Document & body, T & to) {
    memory::MonotonicStack_1k_Resource memoryResource;
    RapidJSONPMRAlloc alloc{&memoryResource};
    StringWriter< T > s{to};
    rapidjson::Writer< StringWriter< T >, rapidjson::UTF8<>, rapidjson::UTF8<>, RapidJSONPMRAlloc > writer{s, &alloc};
    body.Accept(writer);
}

} // namespace rublon

namespace std {
inline auto begin(const rublon::Value & __ils) noexcept {
    return __ils.Begin();
}
inline ::rublon::Value::ConstValueIterator end(const rublon::Value & __ils) noexcept {
    return __ils.End();
}
[[nodiscard]] inline std::size_t size(const rublon::Value & __cont) {
    return __cont.Size();
}

} // namespace std