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

// ScratchArena.hpp
// 一時バッファの高速アリーナアロケータ
// Karatsuba/Toom-Cook の再帰で malloc を回避するために使用
//
// VirtualAlloc ベース: 仮想アドレス空間を大きく予約し、
// 物理メモリはオンデマンドでコミット。リアロケーションが発生しない
// ためポインタが無効化されない。

#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 calx {

class ScratchArena {
    uint64_t* base_ = nullptr;
    size_t offset_ = 0;       // uint64_t 単位の使用済みオフセット
    size_t committed_ = 0;    // uint64_t 単位のコミット済みサイズ
    size_t reserved_ = 0;     // uint64_t 単位の予約サイズ

    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;  // 512KB 単位でコミット

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_) {
            // フォールバック: reserve が失敗した場合は小さめに再試行
            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;

        // 必要に応じて物理メモリをコミット (アドレスは変わらない)
        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; }
};

// thread_local アリーナへのアクセス（static local で ODR 問題を回避）
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 calx