#include <stdexcept>
#include <cassert>
#include <sys/mman.h>
#include "mmap_manager.h"

mmap_manager_t::mmap_manager_t(uint64_t mmap_size)
{
    this->mmap_size = mmap_size;
}

mmap_manager_t::~mmap_manager_t()
{
    for (auto & kv: past_buffers)
    {
        munmap(kv.second.addr, kv.second.size);
    }
    if (active_buffer.addr != NULL)
    {
        munmap(active_buffer.addr, active_buffer.size);
    }
}

void *mmap_manager_t::alloc(uint64_t size)
{
    if (!active_buffer.addr || (active_buffer.pos + size) > active_buffer.size)
    {
        if (active_buffer.addr)
        {
            if (active_buffer.freed >= active_buffer.pos)
                munmap(active_buffer.addr, active_buffer.size);
            else
                past_buffers[active_buffer.addr] = active_buffer;
            active_buffer = { 0 };
        }
        uint64_t new_size = size < mmap_size ? mmap_size : size;
        void *buf = mmap(NULL, new_size, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
        if (!buf)
            throw std::runtime_error(std::string("can't mmap "+std::to_string(new_size)+" bytes"));
        active_buffer = {
            .addr = buf,
            .size = new_size,
            .freed = 0,
            .pos = 0,
        };
    }
    void *res = active_buffer.addr + active_buffer.pos;
    active_buffer.pos += size;
    return res;
}

void mmap_manager_t::free(void *addr, uint64_t size)
{
    auto it = past_buffers.upper_bound(addr);
    if (it != past_buffers.begin())
    {
        if (it == past_buffers.end())
        {
            it--;
            if (addr < it->second.addr || addr >= it->second.addr+it->second.size)
                it = past_buffers.end();
        }
        else
            it--;
    }
    else
        it = past_buffers.end();
    if (it != past_buffers.end())
    {
        assert(addr >= it->second.addr && addr+size <= it->second.addr+it->second.size);
        it->second.freed += size;
        if (it->second.freed >= it->second.pos)
        {
            munmap(it->second.addr, it->second.size);
            past_buffers.erase(it);
        }
    }
    else
    {
        assert(addr < active_buffer.addr+active_buffer.size);
        active_buffer.freed += size;
    }
}