#pragma once
#include <memory>
#include <vector>
#include <stack>

class AlignedMemoryStack {
    char* buffer = nullptr;
    size_t capacity = 0;
    size_t offset = 0;
    std::stack<size_t> markers;

public:
    AlignedMemoryStack(size_t initialSize = 64 * 1024) {
        buffer = new char[initialSize];
        capacity = initialSize;
    }

    ~AlignedMemoryStack() {
        delete[] buffer;
    }

    void* pushStackAlloc(size_t size, size_t align = alignof(std::max_align_t)) {
        size_t aligned = (offset + align - 1) & ~(align - 1);
        if (aligned + size > capacity) {
            // 简单扩容策略：直接翻倍
            size_t newCapacity = (aligned + size) * 2;
            char* newBuffer = new char[newCapacity];
            std::memcpy(newBuffer, buffer, offset);
            delete[] buffer;
            buffer = newBuffer;
            capacity = newCapacity;
        }
        void* ptr = buffer + aligned;
        markers.push(offset);
        offset = aligned + size;
        return ptr;
    }

    template<typename PtrType>
    PtrType* pushStackAllocate(size_t size, size_t align = alignof(std::max_align_t)) {
        return pushStackAlloc(size, align);
    }

    void popStack() {
        offset = markers.top();
        markers.pop();
    }

    void reset() {
        offset = 0;
        markers = std::stack<size_t>();
    }
};