// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#pragma once

#include <algorithm>
#include <cstddef>
#include <mutex>
#include <new>
#include <type_traits>
#include <vector>

// Large blocks that a whole pass's per-image reflection vectors are carved from.
//
// A rotation pass keeps every frame's integrated reflections until its scaling is done - thousands of
// vectors of a few megabytes each, gigabytes together, allocated by the image workers in the
// allocator's per-thread arenas. When the pass hands them back, that memory mostly stays in those
// arenas, as holes between whatever else the workers allocated, and the next pass's workers, being
// new threads, do not reuse it: on a fine-sliced long axis gigabytes of freed reflections were carried
// to the end of the run. Here a block is allocated in one piece (large enough that it is its own
// mapping) and returned in one piece once the last vector in it is gone, so a pass's reflections
// leave nothing behind.
//
// A vector is carved from the current block by bumping a pointer; the space of a vector freed early
// is only reclaimed with its whole block. That suits the use here: the vectors are written once,
// when the frame is integrated, and handed back together.
class ReflectionArena {
public:
    ReflectionArena() = default;
    ReflectionArena(const ReflectionArena &) = delete;
    ReflectionArena &operator=(const ReflectionArena &) = delete;
    ~ReflectionArena() {
        for (auto &b : blocks)
            ::operator delete(b.data);
    }

    void *Allocate(size_t bytes) {
        bytes = (bytes + kAlign - 1) / kAlign * kAlign;
        std::unique_lock ul(m);
        if (blocks.empty() || blocks.back().size - blocks.back().used < bytes) {
            const size_t size = std::max(kBlockBytes, bytes);
            blocks.push_back(Block{static_cast<char *>(::operator new(size)), size, 0, 0});
        }
        Block &b = blocks.back();
        void *p = b.data + b.used;
        b.used += bytes;
        b.live++;
        return p;
    }

    void Deallocate(void *p) {
        std::unique_lock ul(m);
        for (size_t i = 0; i < blocks.size(); ++i) {
            Block &b = blocks[i];
            if (static_cast<char *>(p) < b.data || static_cast<char *>(p) >= b.data + b.size)
                continue;
            if (--b.live == 0) {
                if (i + 1 == blocks.size()) {
                    b.used = 0;   // the block being filled: keep it for the next vector
                } else {
                    ::operator delete(b.data);
                    blocks.erase(blocks.begin() + static_cast<std::ptrdiff_t>(i));
                }
            }
            return;
        }
    }

private:
    // 64 MiB: above the largest size glibc ever serves from its heaps (32 MiB), so every block is a
    // mapping of its own and goes back to the system when freed. A frame's reflections are a few
    // megabytes, so the tail a block cannot fit is a few percent of it.
    static constexpr size_t kBlockBytes = size_t(64) << 20;
    static constexpr size_t kAlign = alignof(std::max_align_t);
    struct Block {
        char *data;
        size_t size, used, live;
    };
    std::mutex m;
    std::vector<Block> blocks;
};

// A std::allocator that carves from a ReflectionArena, or is plain new/delete without one. A copy of
// a container is made on the heap (select_on_container_copy_construction), so only the containers the
// arena's owner fills itself live in the arena; a move carries the arena along with the storage.
template <class T>
struct ArenaAllocator {
    using value_type = T;
    using propagate_on_container_move_assignment = std::true_type;
    using propagate_on_container_swap = std::true_type;

    ReflectionArena *arena = nullptr;

    ArenaAllocator() = default;
    explicit ArenaAllocator(ReflectionArena *a) : arena(a) {}
    template <class U> ArenaAllocator(const ArenaAllocator<U> &o) : arena(o.arena) {}

    T *allocate(size_t n) {
        if (arena)
            return static_cast<T *>(arena->Allocate(n * sizeof(T)));
        return static_cast<T *>(::operator new(n * sizeof(T)));
    }
    void deallocate(T *p, size_t) {
        if (arena)
            arena->Deallocate(p);
        else
            ::operator delete(p);
    }
    ArenaAllocator select_on_container_copy_construction() const { return ArenaAllocator(); }

    template <class U> bool operator==(const ArenaAllocator<U> &o) const { return arena == o.arena; }
    template <class U> bool operator!=(const ArenaAllocator<U> &o) const { return arena != o.arena; }
};