hash_map.hpp

// 哈希表实现
#pragma once
#include <cstring>
#include <cstdint>
#include <memory>
#include <ranges>
#include <utility>
#include <bit>
#if defined(__SSE2__)
#include <immintrin.h>
#endif

namespace misha::details {
    constexpr uint64_t wyhash_secret0 = 0xa0761d6478bd642full;
    constexpr uint64_t wyhash_secret1 = 0xe7037ed1a0b428dbull;

    constexpr inline auto wyhash_mum(uint64_t a, uint64_t b) noexcept -> uint64_t {
        #if defined(__SIZEOF_INT128__)
            auto r = (unsigned __int128)a * b;
            return (r >> 64) ^ r;
        #else
            #error "current compiler/target does not support unsigned __int128"
        #endif
    }

    inline auto fast_range(uint32_t hash, uint32_t capacity) noexcept -> uint32_t {
        return (uint64_t(hash) * capacity) >> 32;
    }

    inline auto match_in_16bytes(const uint8_t* data, uint8_t target) noexcept -> uint16_t {
        #if defined(__SSE2__)
            __m128i bytes = _mm_loadu_si128((const __m128i*)data);
            __m128i match_mask = _mm_set1_epi8(target);
            __m128i cmp = _mm_cmpeq_epi8(bytes, match_mask);
            return _mm_movemask_epi8(cmp);
        #else
            #warning "current compiler/target does not support __SSE2__"
            uint16_t bitmap = 0;
            for (int i = 0; i < 16; i++)
                bitmap |= uint16_t(data[i] == target) << i;
            return bitmap;
        #endif
    }
}


namespace misha {

template <typename T>
struct hash;

template <std::convertible_to<std::string_view> T>
struct hash<T> {
    using is_transparent = void;
    using transparent_type = const std::string_view;

    // Wyhash (v4)
    static auto operator()(const std::string_view sv) noexcept -> uint64_t {
        using namespace misha::details;

        const char* __restrict data = sv.data();
        size_t len = sv.size();
        uint64_t seed = wyhash_secret0 ^ len;

        auto read_bytes = [](const char* __restrict p, size_t n) -> uint64_t {
            uint64_t val = 0;
            return (std::memcpy(&val, p, n), val);
        };

        for (; len > 16; data += 16, len -= 16) {
            uint64_t block0 = read_bytes(data, 8);
            uint64_t block1 = read_bytes(data + 8, 8);
            seed = wyhash_mum(block0 ^ wyhash_secret0, block1 ^ seed);
        }
        uint64_t tail0 = len > 0 ? read_bytes(data, std::min<size_t>(len, 8)) : 0;
        uint64_t tail1 = len > 8 ? read_bytes(data + 8, len - 8) : 0;

        return wyhash_mum(tail0 ^ wyhash_secret0, tail1 ^ seed);

    }
};

template <std::integral T>
struct hash<T> {
    using is_transparent = void;
    using transparent_type = uint64_t;

    constexpr static auto operator()(const T val) noexcept -> uint64_t {
        using namespace misha::details;
        return wyhash_mum(uint64_t(val) ^ wyhash_secret0, wyhash_secret1);
    }
};

template <class Hash, class Key>
concept hashable = requires(Hash hash_func, const Key& key) {
    { hash_func(key) } -> std::convertible_to<uint64_t>;
};

}


namespace misha {

template<class Key, class Value, hashable<Key> Hash = hash<Key>, class KeyEqual = std::equal_to<>>
class hash_map {
public:
    using key_type = Key;
    using mapped_type = Value;
    using size_type = uint32_t;
    struct value_type { key_type key; mapped_type value; };

private:
    template <class Self>
    using value_ref_t = std::conditional_t<std::is_const_v<std::remove_reference_t<Self>>, const Value&, Value&>;

    static constexpr uint8_t CTRL_EMPTY = 0b1000'0000;
    static constexpr uint8_t CTRL_ZOMBIE = 0b1111'1110;

    uint8_t* __restrict data_ = nullptr;
    size_type size_ = 0;
    size_type capacity_ = 0;
    size_type zombies_count_ = 0;
    [[no_unique_address]] Hash hash_func_ = {};
    [[no_unique_address]] KeyEqual is_key_equal_ = {};

    static auto required_size(size_type capacity) noexcept -> size_type {
        return capacity + 15 + alignof(value_type) - 1 + sizeof(value_type) * capacity;
    }

    static auto next_capacity(uint32_t current_capacity, uint32_t min_val = 16) noexcept -> uint32_t {
        if (current_capacity < min_val / 2) return min_val;
        if (current_capacity <= 256) return current_capacity * 2;
        if (current_capacity <= 64 * 1024) return current_capacity * 3 / 2;
        if (current_capacity <= 1024 * 1024) return current_capacity * 7 / 4;
        return (uint64_t)current_capacity * 4 / 3;
    }

    auto calc_hash(const Hash::transparent_type key) const noexcept -> std::pair<uint32_t, uint8_t> {
        const uint64_t h1 = hash_func_(key);
        const uint8_t  h2 = h1 & 0b0111'1111;
        return { h1 >> 7, h2 };
    }
    
    auto ctrls_array() const noexcept -> uint8_t* {
        return data_;
    }

    auto kvs_array() const noexcept -> value_type* {
        return (value_type*)((uintptr_t(data_) + capacity_ + 15 + alignof(value_type) - 1) & ~(alignof(value_type) - 1));
    }

    void set_ctrl(size_type index, uint8_t value) const noexcept {
        ctrls_array()[index] = value;
        if (index < 15)
            ctrls_array()[capacity_ + index] = value;
    }

    void rehash(size_type new_capacity) {
        const uint8_t* __restrict old_ctrls_array = ctrls_array();
        value_type* __restrict old_kvs_array = kvs_array();

        uint8_t* old_data = std::exchange(data_, new uint8_t[required_size(new_capacity)]);
        size_type old_capacity = std::exchange(capacity_, new_capacity);
        zombies_count_ = 0;
        std::memset(ctrls_array(), CTRL_EMPTY, new_capacity + 15);

        for (size_type old_idx = 0; old_idx < old_capacity; old_idx++) {
            if (old_ctrls_array[old_idx] <= 0b0111'1111) {
                auto [h1, h2] = calc_hash(old_kvs_array[old_idx].key);
                
                size_type group_idx = details::fast_range(h1, new_capacity);
                while(true) {
                    uint32_t empty_map = details::match_in_16bytes(ctrls_array() + group_idx, CTRL_EMPTY);
                    if (empty_map) {
                        size_type idx = group_idx + std::countr_zero(empty_map);
                        idx = idx >= capacity_ ? idx - capacity_ : idx;
                        std::construct_at(kvs_array() + idx, std::move(old_kvs_array[old_idx]));
                        std::destroy_at(old_kvs_array + old_idx);
                        set_ctrl(idx, h2);
                        break;
                    }

                    group_idx += 16;
                    group_idx = group_idx >= capacity_ ? group_idx - capacity_ : group_idx;
                }
            }
        }

        delete[] old_data;
    }

    void clean_zombies() {
        // TODO clean zombies in place
        rehash(capacity_);
    }

public:
    hash_map() noexcept = default;

    auto try_emplace(const Hash::transparent_type key, auto&&... value_args) -> std::pair<value_type*, bool> {
        if (size_ + zombies_count_ >= (uint64_t)capacity_ * 7 / 8) [[unlikely]] {
            if (zombies_count_ > (uint64_t)capacity_ * 3 / 16)
                clean_zombies();
            else rehash(next_capacity(capacity_));
        }

        const auto [h1, h2] = calc_hash(key);
        std::optional<size_type> first_zombie_idx = std::nullopt;

        size_type group_idx = details::fast_range(h1, capacity_);
        while(true) {
            const auto get_matched_idx = [&](uint32_t map) -> size_type {
                uint32_t offset = std::countr_zero(map);
                size_t idx = group_idx + offset;
                return idx >= capacity_ ? idx - capacity_ : idx;
            };

            uint32_t h2_map = details::match_in_16bytes(ctrls_array() + group_idx, h2);
            for (; h2_map; h2_map &= h2_map - 1) {
                size_type idx = get_matched_idx(h2_map);
                if (is_key_equal_(kvs_array()[idx].key, key))
                    return { kvs_array() + idx, false };
            }

            if (not first_zombie_idx) {
                uint32_t zombie_map = details::match_in_16bytes(ctrls_array() + group_idx, CTRL_ZOMBIE);
                if (zombie_map)
                    first_zombie_idx = get_matched_idx(zombie_map);
            }

            uint32_t empty_map = details::match_in_16bytes(ctrls_array() + group_idx, CTRL_EMPTY);
            if (empty_map) {
                const auto [insert_idx, is_zombie] = first_zombie_idx ?
                    std::pair { *first_zombie_idx, true } :
                    std::pair { get_matched_idx(empty_map), false };

                if constexpr (sizeof...(value_args) == 1) {
                    if constexpr (std::is_same_v<decltype(value_args...[0]), value_type&&>)
                        std::construct_at(kvs_array() + insert_idx, std::move(value_args...[0]));
                    else
                        std::construct_at(kvs_array() + insert_idx, key, mapped_type(std::forward<decltype(value_args)>(value_args)...));
                } else
                    std::construct_at(kvs_array() + insert_idx, key, mapped_type(std::forward<decltype(value_args)>(value_args)...));

                set_ctrl(insert_idx, h2);

                size_++;
                zombies_count_ = is_zombie ? zombies_count_ - 1 : zombies_count_;

                return { kvs_array() + insert_idx, true };
            }

            group_idx += 16;
            group_idx = group_idx >= capacity_ ? group_idx - capacity_ : group_idx;
        }
    }

    void insert_or_assign(value_type&& val) {
        auto [ptr, inserted] = try_emplace(val.key, std::move(val));
        if (not inserted)
            *ptr = std::move(val);
    }

    auto find(this auto&& self, const Hash::transparent_type key) noexcept -> std::optional<value_ref_t<decltype(self)>> {
        if (self.size_ == 0) [[unlikely]] return std::nullopt;

        const auto [h1, h2] = self.calc_hash(key);
        size_type group_idx = details::fast_range(h1, self.capacity_);
        while(true) {
            uint32_t h2_map = details::match_in_16bytes(self.ctrls_array() + group_idx, h2);
            for (; h2_map; h2_map &= h2_map - 1) {
                size_t idx = group_idx + std::countr_zero(h2_map);
                idx = idx >= self.capacity_ ? idx - self.capacity_ : idx;

                if (self.is_key_equal_(self.kvs_array()[idx].key, key))
                    return self.kvs_array()[idx].value;
            }

            uint32_t empty_map = details::match_in_16bytes(self.ctrls_array() + group_idx, CTRL_EMPTY);
            if (empty_map)
                return std::nullopt;

            group_idx += 16;
            group_idx = group_idx >= self.capacity_ ? group_idx - self.capacity_ : group_idx;
        }
    }

    void erase(const Hash::transparent_type key) noexcept {
        if (auto val = find(key))
            erase_at(*val);
    }

    void erase_at(const mapped_type& value) noexcept {
        if (uintptr_t(&value) >= (uintptr_t)kvs_array() and uintptr_t(&value) < uintptr_t(kvs_array() + capacity_)) {
            size_type idx = (uintptr_t(&value) - uintptr_t(kvs_array())) / sizeof(value_type);
            std::destroy_at(kvs_array() + idx);
            set_ctrl(idx, CTRL_ZOMBIE);
            size_--;
            zombies_count_++;
        }
    }

    auto entries(this auto&& self) noexcept {
        return
            std::views::iota(size_type(0), self.capacity_) |
            std::views::filter([ctrls = self.ctrls_array()](size_type i) { return ctrls[i] <= 0b0111'1111; }) |
            std::views::transform([kvs = self.kvs_array()](size_type i) -> std::pair<Key&, value_ref_t<decltype(self)>> {
                return { kvs[i].key, kvs[i].value };
            });
    }

    void clear() noexcept {
        if (capacity_ == 0) [[unlikely]] return;

        for (size_type idx = 0; idx < capacity_; idx++) {
            if (ctrls_array()[idx] <= 0b0111'1111)
                std::destroy_at(kvs_array() + idx);
            ctrls_array()[idx] = CTRL_EMPTY;
        }

        std::memset(ctrls_array() + capacity_, CTRL_EMPTY, 15);
        size_ = zombies_count_ = 0;
    }

    auto contains(const Hash::transparent_type key) const noexcept -> bool { return find(key).has_value(); }
    auto size() const noexcept -> size_type { return size_; }
    auto capacity() const noexcept -> size_type { return capacity_; }
    auto empty() const noexcept -> bool { return size_ == 0; }

    auto operator[](const Hash::transparent_type key) -> mapped_type& {
        return try_emplace(key).first->value;
    }
    
    hash_map(hash_map&& other) noexcept :
        data_(std::exchange(other.data_, nullptr)),
        size_(std::exchange(other.size_, 0)),
        capacity_(std::exchange(other.capacity_, 0)),
        zombies_count_(std::exchange(other.zombies_count_, 0)),
        hash_func_(std::move(other.hash_func_)),
        is_key_equal_(std::move(other.is_key_equal_))
    {}

    ~hash_map() noexcept {
        clear();
        delete[] data_;
    }
};

}