// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// ScratchArena.hpp
// Fast arena allocator for temporary buffers
// Used to avoid malloc inside Karatsuba/Toom-Cook recursion
//
// VirtualAlloc-based: reserve a large virtual address range up front and
// commit physical memory on demand. Reallocation never occurs, so existing
// pointers never get invalidated.

#pragma once

#include <cstdint>
#include <cstddef>
#include <stdexcept>

#ifdef _WIN32
#ifndef NOMINMAX
#define NOMINMAX
#endif
#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif
#include <windows.h>
#else
#include <sys/mman.h>
#include <unistd.h>
#endif

namespace sangi {

class ScratchArena {
    uint64_t* base_ = nullptr;
    size_t offset_ = 0;       // used offset in uint64_t units
    size_t committed_ = 0;    // committed size in uint64_t units
    size_t reserved_ = 0;     // reserved size in uint64_t units

    static constexpr size_t RESERVE_WORDS = size_t(1) << 29;  // 4GB (512M words * 8 bytes)
    static constexpr size_t COMMIT_GRANULARITY = size_t(1) << 16;  // commit in 512KB chunks

public:
    ScratchArena() {
        reserved_ = RESERVE_WORDS;
        size_t bytes = reserved_ * sizeof(uint64_t);
#ifdef _WIN32
        base_ = static_cast<uint64_t*>(
            VirtualAlloc(nullptr, bytes, MEM_RESERVE, PAGE_NOACCESS));
#else
        base_ = static_cast<uint64_t*>(
            mmap(nullptr, bytes, PROT_NONE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));
        if (base_ == MAP_FAILED) base_ = nullptr;
#endif
        if (!base_) {
            // Fallback: if reserve fails, retry with a smaller size
            reserved_ = size_t(1) << 24;  // 128MB
            bytes = reserved_ * sizeof(uint64_t);
#ifdef _WIN32
            base_ = static_cast<uint64_t*>(
                VirtualAlloc(nullptr, bytes, MEM_RESERVE, PAGE_NOACCESS));
#else
            base_ = static_cast<uint64_t*>(
                mmap(nullptr, bytes, PROT_NONE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));
            if (base_ == MAP_FAILED) base_ = nullptr;
#endif
        }
    }

    ~ScratchArena() {
        if (base_) {
#ifdef _WIN32
            VirtualFree(base_, 0, MEM_RELEASE);
#else
            munmap(base_, reserved_ * sizeof(uint64_t));
#endif
        }
    }

    ScratchArena(const ScratchArena&) = delete;
    ScratchArena& operator=(const ScratchArena&) = delete;

    uint64_t* alloc_limbs(size_t n) {
        if (n == 0) return nullptr;
        // 8-word alignment (64 bytes, cache line)
        size_t aligned_n = (n + 7) & ~size_t(7);
        size_t new_end = offset_ + aligned_n;

        // Commit physical memory on demand (the address does not change)
        if (new_end > committed_) {
            size_t need = new_end - committed_;
            size_t commit = ((need + COMMIT_GRANULARITY - 1) / COMMIT_GRANULARITY)
                            * COMMIT_GRANULARITY;
            if (committed_ + commit > reserved_) {
                commit = reserved_ - committed_;
                if (new_end > reserved_) {
                    throw std::runtime_error("ScratchArena: exceeded reserved address space");
                }
            }
#ifdef _WIN32
            void* p = VirtualAlloc(base_ + committed_,
                                   commit * sizeof(uint64_t),
                                   MEM_COMMIT, PAGE_READWRITE);
            if (!p) throw std::runtime_error("ScratchArena: VirtualAlloc commit failed");
#else
            int ret = mprotect(base_ + committed_,
                               commit * sizeof(uint64_t),
                               PROT_READ | PROT_WRITE);
            if (ret != 0) throw std::runtime_error("ScratchArena: mprotect failed");
#endif
            committed_ += commit;
        }

        uint64_t* p = base_ + offset_;
        offset_ = new_end;
        return p;
    }

    size_t mark() const { return offset_; }
    void rewind(size_t saved) { offset_ = saved; }
};

// Access to the thread_local arena (a static local avoids ODR issues)
inline ScratchArena& getThreadArena() {
    static thread_local ScratchArena arena;
    return arena;
}

class ScratchScope {
    size_t saved_;
public:
    ScratchScope() : saved_(getThreadArena().mark()) {}
    ~ScratchScope() { getThreadArena().rewind(saved_); }
    ScratchScope(const ScratchScope&) = delete;
    ScratchScope& operator=(const ScratchScope&) = delete;
};

} // namespace sangi