AronaCore/core/misc/mempool.h

#pragma once
#include <memory>
#include <queue>
#include <vector>
#include <bits/ranges_algo.h>

class memory_pool {
public:
    memory_pool(const size_t chunk_size, const size_t chunks_per_block)
        : chunk_size_(chunk_size), chunks_per_block_(chunks_per_block), current_block_(nullptr), free_list_(nullptr) {
        allocate_block();
    }
    memory_pool(memory_pool&& other) noexcept
        : chunk_size_(other.chunk_size_), chunks_per_block_(other.chunks_per_block_), blocks_(std::move(other.blocks_)),
          current_block_(other.current_block_), free_list_(other.free_list_) {
        other.current_block_ = nullptr;
        other.free_list_ = nullptr;
    }

    ~memory_pool() {
        for (void* block : blocks_) {
            delete[] static_cast<char*>(block);
        }
    }

    auto allocate() -> void* {
        if (!free_list_) {
            allocate_block();
        }

        std::lock_guard lock(mutex_);
        void* chunk = free_list_;
        free_list_ = *static_cast<void**>(free_list_);
        return chunk;
    }

    void deallocate(void* chunk) {
        std::lock_guard lock(mutex_);
        *static_cast<void**>(chunk) = free_list_;
        free_list_ = chunk;
    }

private:
    void allocate_block() {
        const size_t block_size = chunk_size_ * chunks_per_block_;
        const auto new_block = new char[block_size];
        {
            std::lock_guard lock(mutex_);
            blocks_.push_back(new_block);
            current_block_ = new_block;
        }

        for (size_t i = 0; i < chunks_per_block_; ++i) {
            void* chunk = new_block + i * chunk_size_;
            deallocate(chunk);
        }
    }

    size_t chunk_size_;
    size_t chunks_per_block_;
    std::vector<void*> blocks_;
    void* current_block_;
    void* free_list_;
    std::mutex mutex_;
};

template<class T>
class obj_mempool {
public:
    template<typename ...Args>
    static auto construct(Args&&... args) -> T* {
        auto obj = alloc();
        new (obj) T(std::forward<Args>(args)...);
        return obj;
    }
    static auto construct() -> T* {
        auto obj = alloc();
        new (obj) T();
        return obj;
    }
    static void free(T* p) {
        p->~T();
        deallocate(p);
    }
    static void free_all() {
        for (auto obj : objs_) {
            obj->~T();
            pool_.deallocate(obj);
        }
        objs_.clear();
    }
    static auto objs() -> const std::vector<T*>& {
        return objs_;
    }
    static auto has_obj(T* p) -> bool {
        return std::find(objs_.begin(), objs_.end(), p) != objs_.end();
    }
    static auto safe_free(T* p) -> bool {
        if (has_obj(p)) {
            free(p);
            return true;
        }
        return false;
    }

    static auto alloc() -> T* {
        T* p = static_cast<T*>(pool_.allocate());
        if (free_indices_.empty()) {
            objs_.push_back(p);
        } else {
            size_t index = free_indices_.front();
            free_indices_.pop();
            objs_.insert(objs_.begin() + index, p);
        }
        return p;
    }
    static void deallocate(T* p) {
        pool_.deallocate(p);
        // find p index
        auto it = std::find(objs_.begin(), objs_.end(), p);
        if (it != objs_.end()) {
            size_t index = std::distance(objs_.begin(), it);
            free_indices_.push(index);
        }
        // remove p from objs_
        objs_.erase(it);
    }


	static auto begin()
	{
		return objs_.begin();
	}
	static auto end()
	{
		return objs_.end();
	}
	static auto size()
	{
		return objs_.size();
	}
	T* operator[](size_t index)
	{
		return objs_[index];
	}
private:
    inline static auto pool_ = memory_pool(sizeof(T), 64);
    inline static std::vector<T*> objs_;
    inline static std::queue<size_t> free_indices_;
};

template<class T>
class pool_obj {
public:
    auto operator new(size_t size) -> void* {
        return obj_mempool<T>::alloc();
    }
    void operator delete(void* p) {
        obj_mempool<T>::deallocate(static_cast<T*>(p));
    }
};

template<class T>
T* get_pool_obj()
{
	return obj_mempool<T>::construct();
}

template<class T, typename ...Args>
T* get_pool_obj(Args&&... InArgs)
{
	return obj_mempool<T>::construct(std::forward<Args>(InArgs)...);
}

template<class T>
void free_pool_obj(T* P)
{
	obj_mempool<T>::free(P);
}

// 基础模板，用于递归计算最大类大小
template<class First, class... Rest>
struct max_class_size {
	static constexpr size_t value = sizeof(First) > max_class_size<Rest...>::value ? sizeof(First) : max_class_size<Rest...>::value;
};

// 递归终止条件
template<class Last>
struct max_class_size<Last> {
	static constexpr size_t value = sizeof(Last);
};

template<class BaseClass, class ...ChildClasses>
class inherit_obj_pool
{
	static constexpr size_t MAX_CHILD_SIZE = max_class_size<ChildClasses...>::value;
public:
	template<class T, typename ...Args>
	static auto construct(Args&&... InArgs) -> T*
	{
		static_assert((std::is_same_v<T, ChildClasses> || ...), "T必须是ChildClasses中的一个");
		auto obj = alloc<T>();
		new (obj) T(std::forward<Args>(InArgs)...);
		return obj;
	}
	template<class T>
	static auto construct() -> T*
	{
		static_assert((std::is_same_v<T, ChildClasses> || ...), "T必须是ChildClasses中的一个");
		auto obj = alloc<T>();
		new (obj) T();
		return obj;
	}
	static void free(BaseClass* p)
	{
		if (!p)
			return;
		p->~BaseClass();
		deallocate(p);
	}
	static void free_all()
	{
		for (auto obj : objs_)
		{
			obj->~BaseClass();
			pool_.deallocate(obj);
		}
		objs_.clear();
	}
	static auto get_objs() -> const std::vector<BaseClass*>&
	{
		return objs_;
	}
	static auto has_obj(BaseClass* P) -> bool
	{
		return std::find(objs_.begin(), objs_.end(), P) != objs_.end();
	}
	static auto safe_free(BaseClass* P) -> bool
	{
		if (has_obj(P))
		{
			free(P);
			return true;
		}
		return false;
	}

	template<class T>
	static auto alloc() -> T*
	{
		T* P = static_cast<T*>(pool_.allocate());
		if (free_indexs_.empty())
		{
			objs_.push_back(P);
		}
		else
		{
			size_t index = free_indexs_.front();
			free_indexs_.pop();
			objs_.insert(objs_.begin() + index, P);
		}
		return P;
	}
	static void deallocate(BaseClass* P)
	{
		pool_.deallocate(P);
		auto it = std::find(objs_.begin(), objs_.end(), P);
		if (it != objs_.end())
		{
			const size_t index = std::distance(objs_.begin(), it);
			free_indexs_.push(index);
		}
		objs_.erase(it);
	}

	static auto begin()
	{
		return objs_.begin();
	}
	static auto end()
	{
		return objs_.end();
	}
	static auto size()
	{
		return objs_.size();
	}
	BaseClass* operator[](size_t index)
	{
		return objs_[index];
	}
private:
	inline static auto pool_ = memory_pool(MAX_CHILD_SIZE, 128);
	inline static std::vector<BaseClass*> objs_;
	inline static std::queue<size_t> free_indexs_;
};