XCEngine/engine/include/XCEngine/Containers/HashMap.h

#pragma once

#include <cstddef>
#include <utility>
#include <functional>
#include "../Memory/Allocator.h"
#include "Array.h"

namespace XCEngine {
namespace Containers {

template<typename Key, typename Value>
class HashMap {
public:
    struct Pair {
        Key first;
        Value second;
    };

    HashMap() {
        m_buckets.Resize(DefaultBucketCount);
    }
    explicit HashMap(size_t bucketCount, Memory::IAllocator* allocator = nullptr);
    ~HashMap();

    HashMap(const HashMap& other);
    HashMap(HashMap&& other) noexcept;
    HashMap& operator=(const HashMap& other);
    HashMap& operator=(HashMap&& other) noexcept;

    Value& operator[](const Key& key);
    Value* Find(const Key& key);
    const Value* Find(const Key& key) const;
    bool Contains(const Key& key) const;

    bool Insert(const Key& key, const Value& value);
    bool Insert(const Key& key, Value&& value);
    bool Insert(Pair&& pair);
    bool Erase(const Key& key);
    void Clear();

    size_t Size() const { return m_size; }
    bool Empty() const { return m_size == 0; }

    void SetAllocator(Memory::IAllocator* allocator) { m_allocator = allocator; }

private:
    struct Bucket {
        Array<Pair> pairs;
    };

    static constexpr size_t DefaultBucketCount = 16;

    Array<Bucket> m_buckets;
    size_t m_bucketCount = DefaultBucketCount;
    size_t m_size = 0;
    float m_loadFactor = 0.75f;
    Memory::IAllocator* m_allocator = nullptr;

    size_t GetBucketIndex(const Key& key) const;
    void Resize();
    typename Array<Pair>::Iterator FindInBucket(Bucket& bucket, const Key& key);
    typename Array<Pair>::ConstIterator FindInBucket(const Bucket& bucket, const Key& key) const;
};

template<typename Key, typename Value>
HashMap<Key, Value>::HashMap(size_t bucketCount, Memory::IAllocator* allocator)
    : m_bucketCount(bucketCount), m_allocator(allocator) {
    if (m_bucketCount == 0) {
        m_bucketCount = 16;
    }
    m_buckets.Resize(m_bucketCount);
}

template<typename Key, typename Value>
HashMap<Key, Value>::~HashMap() {
    Clear();
}

template<typename Key, typename Value>
HashMap<Key, Value>::HashMap(const HashMap& other)
    : m_bucketCount(other.m_bucketCount), m_size(other.m_size), m_loadFactor(other.m_loadFactor), m_allocator(other.m_allocator) {
    m_buckets.Resize(m_bucketCount);
    for (size_t i = 0; i < m_bucketCount; ++i) {
        m_buckets[i].pairs = other.m_buckets[i].pairs;
    }
}

template<typename Key, typename Value>
HashMap<Key, Value>::HashMap(HashMap&& other) noexcept
    : m_buckets(std::move(other.m_buckets)),
      m_bucketCount(other.m_bucketCount),
      m_size(other.m_size),
      m_loadFactor(other.m_loadFactor),
      m_allocator(other.m_allocator) {
    other.m_buckets.Clear();
    other.m_bucketCount = 0;
    other.m_size = 0;
}

template<typename Key, typename Value>
HashMap<Key, Value>& HashMap<Key, Value>::operator=(const HashMap& other) {
    if (this != &other) {
        Clear();
        m_bucketCount = other.m_bucketCount;
        m_size = other.m_size;
        m_loadFactor = other.m_loadFactor;
        m_allocator = other.m_allocator;
        m_buckets.Resize(m_bucketCount);
        for (size_t i = 0; i < m_bucketCount; ++i) {
            m_buckets[i].pairs = other.m_buckets[i].pairs;
        }
    }
    return *this;
}

template<typename Key, typename Value>
HashMap<Key, Value>& HashMap<Key, Value>::operator=(HashMap&& other) noexcept {
    if (this != &other) {
        Clear();
        m_buckets = std::move(other.m_buckets);
        m_bucketCount = other.m_bucketCount;
        m_size = other.m_size;
        m_loadFactor = other.m_loadFactor;
        m_allocator = other.m_allocator;

        other.m_buckets.Clear();
        other.m_bucketCount = 0;
        other.m_size = 0;
    }
    return *this;
}

template<typename Key, typename Value>
Value& HashMap<Key, Value>::operator[](const Key& key) {
    size_t bucketIndex = GetBucketIndex(key);
    Bucket& bucket = m_buckets[bucketIndex];

    auto it = FindInBucket(bucket, key);
    if (it != bucket.pairs.end()) {
        return it->second;
    }

    bucket.pairs.EmplaceBack(Pair{key, Value()});
    ++m_size;

    if (static_cast<float>(m_size) / static_cast<float>(m_bucketCount) > m_loadFactor) {
        Resize();
    }

    return bucket.pairs.Back().second;
}

template<typename Key, typename Value>
Value* HashMap<Key, Value>::Find(const Key& key) {
    size_t bucketIndex = GetBucketIndex(key);
    Bucket& bucket = m_buckets[bucketIndex];

    auto it = FindInBucket(bucket, key);
    if (it != bucket.pairs.end()) {
        return &it->second;
    }
    return nullptr;
}

template<typename Key, typename Value>
const Value* HashMap<Key, Value>::Find(const Key& key) const {
    size_t bucketIndex = GetBucketIndex(key);
    const Bucket& bucket = m_buckets[bucketIndex];

    auto it = FindInBucket(bucket, key);
    if (it != bucket.pairs.end()) {
        return &it->second;
    }
    return nullptr;
}

template<typename Key, typename Value>
bool HashMap<Key, Value>::Contains(const Key& key) const {
    size_t bucketIndex = GetBucketIndex(key);
    const Bucket& bucket = m_buckets[bucketIndex];
    return FindInBucket(bucket, key) != bucket.pairs.end();
}

template<typename Key, typename Value>
bool HashMap<Key, Value>::Insert(const Key& key, const Value& value) {
    size_t bucketIndex = GetBucketIndex(key);
    Bucket& bucket = m_buckets[bucketIndex];

    auto it = FindInBucket(bucket, key);
    if (it != bucket.pairs.end()) {
        it->second = value;
        return false;
    }

    bucket.pairs.EmplaceBack(Pair{key, value});
    ++m_size;

    if (static_cast<float>(m_size) / static_cast<float>(m_bucketCount) > m_loadFactor) {
        Resize();
    }

    return true;
}

template<typename Key, typename Value>
bool HashMap<Key, Value>::Insert(const Key& key, Value&& value) {
    size_t bucketIndex = GetBucketIndex(key);
    Bucket& bucket = m_buckets[bucketIndex];

    auto it = FindInBucket(bucket, key);
    if (it != bucket.pairs.end()) {
        it->second = std::move(value);
        return false;
    }

    bucket.pairs.EmplaceBack(Pair{key, std::move(value)});
    ++m_size;

    if (static_cast<float>(m_size) / static_cast<float>(m_bucketCount) > m_loadFactor) {
        Resize();
    }

    return true;
}

template<typename Key, typename Value>
bool HashMap<Key, Value>::Insert(Pair&& pair) {
    return Insert(std::move(pair.first), std::move(pair.second));
}

template<typename Key, typename Value>
bool HashMap<Key, Value>::Erase(const Key& key) {
    size_t bucketIndex = GetBucketIndex(key);
    Bucket& bucket = m_buckets[bucketIndex];

    auto it = FindInBucket(bucket, key);
    if (it != bucket.pairs.end()) {
        size_t index = it - bucket.pairs.begin();
        
        if (index != bucket.pairs.Size() - 1) {
            bucket.pairs[index] = std::move(bucket.pairs.Back());
        }
        bucket.pairs.PopBack();
        --m_size;
        return true;
    }
    return false;
}

template<typename Key, typename Value>
void HashMap<Key, Value>::Clear() {
    for (size_t i = 0; i < m_buckets.Size(); ++i) {
        m_buckets[i].pairs.Clear();
    }
    m_size = 0;
}

template<typename Key, typename Value>
size_t HashMap<Key, Value>::GetBucketIndex(const Key& key) const {
    if (m_bucketCount == 0) {
        return 0;
    }
    std::hash<Key> hasher;
    return hasher(key) % m_bucketCount;
}

template<typename Key, typename Value>
void HashMap<Key, Value>::Resize() {
    size_t newBucketCount = m_bucketCount * 2;
    Array<Bucket> newBuckets;
    newBuckets.Resize(newBucketCount);

    for (size_t i = 0; i < m_bucketCount; ++i) {
        Bucket& oldBucket = m_buckets[i];
        for (size_t j = 0; j < oldBucket.pairs.Size(); ++j) {
            Pair& pair = oldBucket.pairs[j];
            std::hash<Key> hasher;
            size_t newIndex = hasher(pair.first) % newBucketCount;
            newBuckets[newIndex].pairs.EmplaceBack(std::move(pair));
        }
    }

    m_buckets = std::move(newBuckets);
    m_bucketCount = newBucketCount;
}

template<typename Key, typename Value>
typename Array<typename HashMap<Key, Value>::Pair>::Iterator HashMap<Key, Value>::FindInBucket(Bucket& bucket, const Key& key) {
    for (auto it = bucket.pairs.begin(); it != bucket.pairs.end(); ++it) {
        if (it->first == key) {
            return it;
        }
    }
    return bucket.pairs.end();
}

template<typename Key, typename Value>
typename Array<typename HashMap<Key, Value>::Pair>::ConstIterator HashMap<Key, Value>::FindInBucket(const Bucket& bucket, const Key& key) const {
    for (auto it = bucket.pairs.begin(); it != bucket.pairs.end(); ++it) {
        if (it->first == key) {
            return it;
        }
    }
    return bucket.pairs.end();
}

} // namespace Containers
} // namespace XCEngine