#include "shared.h"

#include <util/system/guard.h>

#include <sys/mman.h>
#include <errno.h>

namespace NSrvKernel {
    size_t TSharedAllocator::TypeIdCounter = 0;


    TSharedAllocator::TAllocatedChunk::TAllocatedChunk(void* addr, void* data, size_t size) noexcept
        : Address(addr)
        , Data(data)
        , MappingSize(size)
    {
        Y_VERIFY(addr && size, "a broken invarant");
    }

    TSharedAllocator::TAllocatedChunk::~TAllocatedChunk() {
        munmap(Address, MappingSize);
    }


    TSharedAllocator::TSharedAllocator(size_t minAllocationChunkSize)
        : PageSize_(NSystemInfo::GetPageSize())
        , NextChunkSize_(Max(minAllocationChunkSize, 4 * PageSize_)) // 2 pages + 2 guard pages
    {}

    void* TSharedAllocator::AllocateInChunk(TSharedAllocator::TAllocatedChunk& chunk, size_t size, size_t alignment) {
        if (chunk.MappingSize - chunk.AllocatedSize < size) {
            return nullptr;
        }

        ui8* chunkStart = reinterpret_cast<ui8*>(chunk.Data);
        ui8* objectAddr = AlignUp(chunkStart + chunk.AllocatedSize, alignment);

        if (objectAddr + size > chunkStart + chunk.MappingSize) {
            return nullptr;
        }

        chunk.AllocatedSize = objectAddr + size - chunkStart;

        return objectAddr;
    }

    TSharedAllocator::TAllocatedChunk& TSharedAllocator::AllocateNewChunk(TDeque<TAllocatedChunk>& allocatedChunks, size_t size) {
        size = AlignUp(size, NextChunkSize_);
        // Exp growing of chunks reduces memory fragmentation on guard pages
        NextChunkSize_ *= 2;
        void* addr = mmap(nullptr, size + 2 * PageSize_, PROT_NONE, MAP_SHARED | MAP_ANONYMOUS, -1, 0);

        if (addr == MAP_FAILED) {
            ythrow TSystemError(errno) << "cannot allocate shared chunk";
        }

        void* dataPtr = reinterpret_cast<uint8_t*>(addr) + PageSize_;
        int res = mprotect(dataPtr, size, PROT_READ | PROT_WRITE);
        if (res != 0) {
            ythrow TSystemError(errno) << "cannot unprotect data pages";
        }

        allocatedChunks.emplace_back(addr, dataPtr, size);
        return allocatedChunks.back();
    }

    void* TSharedAllocator::AllocateImpl(size_t typeId, size_t size, size_t alignment) {
        Y_ENSURE(!Frozen_, "shared allocator is frozen (trying to allocate shared memory in child?)");

        void* objectAddr = nullptr;

        with_lock (Lock_) {
            TDeque<TAllocatedChunk>& allocatedChunks = AllocatedChunks_[typeId];

            if (allocatedChunks.size() > 0) {
                objectAddr = AllocateInChunk(allocatedChunks.back(), size, alignment);
            }

            if (!objectAddr) {
                TAllocatedChunk& chunk = AllocateNewChunk(allocatedChunks, size);
                objectAddr = AllocateInChunk(chunk, size, alignment);
                Y_VERIFY(!!objectAddr);
            }

            return objectAddr;
        }
    }

    size_t TSharedAllocator::MappedBytes() const noexcept {
        with_lock (Lock_) {
            size_t mapped = 0;
            for (const auto& entry : AllocatedChunks_) {
                for (const auto& chunk : entry.second) {
                    mapped += chunk.MappingSize;
                }
            }
            return mapped;
        }
    }

    size_t TSharedAllocator::AllocatedBytes() const noexcept {
        with_lock (Lock_) {
            size_t allocated = 0;
            for (const auto& entry : AllocatedChunks_) {
                for (const auto& chunk : entry.second) {
                    allocated += chunk.AllocatedSize;
                }
            }
            return allocated;
        }
    }
}
