ImplicitSurfaceNetwork/shared_module/container/hashed_refcount_hive.hpp

#pragma once

#include "hive.hpp"
#include "hashmap.hpp"

namespace detail
{
template <typename K>
struct hasher {
    size_t operator()(const K& k) const;
};

// template <typename K>
// struct eq_compare
// {
//     bool operator()(const K& lhs, const K& rhs) const;
// };

template <typename K, typename V>
struct default_elem_ctor {
    V operator()(const K& k) const;
};

template <typename K, typename V, typename Allocator = std::allocator<V>>
struct hashed_refcount_hive {
    using iterator = typename hive<V, Allocator>::iterator;

public:
    std::pair<iterator, bool> acquire(const K& key)
    {
        size_t hash_key = hasher<K>()(key);
        auto   iter     = refcount_data_map.find(hash_key);
        if (iter == refcount_data_map.end()) {
            auto data_iter = data.emplace(default_elem_ctor<K, V>{}(key));
            refcount_data_map.emplace(hash_key, std::make_pair(1, data_iter));
            return {data_iter, true};
        } else {
            iter->second.first += 1;
            return {iter->second.second, false};
        }
    }

    void release(const K& key)
    {
        size_t hash_key = hasher<K>()(key);
        auto   iter     = refcount_data_map.find(hash_key);
        if (iter == refcount_data_map.end()) throw std::runtime_error("Key not found in refcount map.");
        if (--iter->second.first == 0) {
            data.erase(iter->second.second);
            refcount_data_map.erase(iter);
        }
    }

    void release(const iterator& ptr)
    {
        for (auto iter = refcount_data_map.begin(); iter != refcount_data_map.end(); ++iter) {
            if (iter->second.second == ptr) {
                if (--iter->second.first == 0) {
                    data.erase(iter->second.second);
                    refcount_data_map.erase(iter);
                }
                return;
            }
        }
        throw std::runtime_error("Pointer not found in refcount map.");
    }

protected:
    hive<V, Allocator>                                 data{};
    // manually calculate hash for the key, so that we do not need to store the key
    flat_hash_map<size_t, std::pair<size_t, iterator>> refcount_data_map{};
};

template <typename T, typename Tag>
struct tagged_hasher;

template <typename K, typename V, typename Tag, typename Allocator = std::allocator<V>>
struct tagged_hashed_refcount_hive {
    using iterator = typename hive<V, Allocator>::iterator;

public:
    std::pair<iterator, bool> acquire(const K& key)
    {
        size_t hash_key = tagged_hasher<K, Tag>()(key);
        auto   iter     = refcount_data_map.find(hash_key);
        if (iter == refcount_data_map.end()) {
            auto data_iter = data.emplace(default_elem_ctor<K, V>{}(key));
            refcount_data_map.emplace(hash_key, std::make_pair(1, data_iter));
            return {data_iter, true};
        } else {
            iter->second.first += 1;
            return {iter->second.second, false};
        }
    }

    void release(const K& key)
    {
        size_t hash_key = tagged_hasher<K, Tag>()(key);
        auto   iter     = refcount_data_map.find(hash_key);
        if (iter == refcount_data_map.end()) throw std::runtime_error("Key not found in refcount map.");
        if (--iter->second.first == 0) {
            data.erase(iter->second.second);
            refcount_data_map.erase(iter);
        }
    }

    void release(const iterator& ptr)
    {
        for (auto iter = refcount_data_map.begin(); iter != refcount_data_map.end(); ++iter) {
            if (iter->second.second == ptr) {
                if (--iter->second.first == 0) {
                    data.erase(iter->second.second);
                    refcount_data_map.erase(iter);
                }
                return;
            }
        }
        throw std::runtime_error("Pointer not found in refcount map.");
    }

protected:
    hive<V, Allocator>                                 data{};
    // manually calculate hash for the key, so that we do not need to store the key
    flat_hash_map<size_t, std::pair<size_t, iterator>> refcount_data_map{};
};
} // namespace detail

template <typename K, typename V>
using hashed_refcount_hive_mp = detail::hashed_refcount_hive<K, V, mi_stl_allocator<V>>;

template <typename K, typename V, typename Tag>
using tagged_hashed_refcount_hive_mp = detail::tagged_hashed_refcount_hive<K, V, Tag, mi_stl_allocator<V>>;