HashMap.h

#pragma once

#include "HashMapIterators.h"

namespace SE {

template<typename TKey, typename TValue, template<typename> typename THasherTemplate>
class THashMap {
private:
    using FEntry = Impl::THashMapEntry<TKey, TValue>;
    using THasher = THasherTemplate<TKey>;

public:
    using FAllocator = TBlockAllocator<sizeof(FEntry), (uint32)alignof(Impl::THashMapEntry<TKey, TValue>)>;
    using FIterator = THashMapIterator<TKey, TValue, FAllocator>;
    using FConstIterator = THashMapConstIterator<TKey, TValue>;

    THashMap() {
        Entries.Add((FEntry *)1);
    }

    THashMap(const THashMap &Other) {
        Construct(Other);
    }

    THashMap &operator=(const THashMap &Other) {
        Clear();
        Entries.Last() = nullptr;
        Construct(Other);
        return *this;
    }

    THashMap(THashMap &&Other) {
        Entries.Add((FEntry *)1);
        MoveConstruct(Move(Other));
    }

    THashMap &operator=(THashMap &&Other) {
        Clear();
        Entries.Resize(1);
        Entries[0] = (FEntry *)1;
        MoveConstruct(Move(Other));
        return *this;
    }

    ~THashMap() {
        Clear();
    }

    // Invalidated by external container mutation.
    FConstIterator CreateConstIterator() const {
        FEntry *const *End = &Entries.Last();
        for (FEntry *const *Entry = &Entries.First(); Entry != End; ++Entry) {
            if (*Entry)
                return {*Entry, Entry};
        }
        return {*End, End};
    }

    FConstIterator begin() const = delete;
    FConstIterator end() const = delete;

    // Invalidated by external container mutation.
    FIterator CreateIterator() {
        FEntry **End = &Entries.Last();
        for (FEntry **Entry = &Entries.First(); Entry != End; ++Entry) {
            if (*Entry)
                return {&ElementCount, *Entry, Entry, &Allocator};
        }
        return {&ElementCount, *End, End, &Allocator};
    }

    FIterator begin() = delete;
    FIterator end() = delete;

    bool Contains(const TKey &Key) const {
        auto BucketIndex = ComputeBucket(Key);
        return FindEntry(Key, BucketIndex);
    }

    bool Contains(const TKey &Key, uint32 BucketIndex) const {
        Check(ComputeBucket(Key) == BucketIndex);
        return FindEntry(Key, BucketIndex);
    }

    uint32 ComputeBucket(const TKey &Key) const {
        return BucketFromHash(THasher::Get(Key));
    }

    TValue &Create(const TKey &Key) {
        Check(!Contains(Key));
        if (IsFull()) {
            Grow();
        }
        auto BucketIndex = ComputeBucket(Key);
        return CreateImpl(Key, BucketIndex);
    }
    
    TValue &Create(const TKey &Key, uint32 BucketIndex) {
        Check(BucketIndex == ComputeBucket(Key));
        Check(!Contains(Key, BucketIndex));
        if (IsFull()) {
            Grow();
            BucketIndex = ComputeBucket(Key);
        }
        return CreateImpl(Key, BucketIndex);
    }

    TValue& Add(const TKey &Key, const TValue &Value) {
        Check(!Contains(Key));
        if (IsFull()) {
            Grow();
        }
        auto BucketIndex = ComputeBucket(Key);
        return AddImpl(Key, Value, BucketIndex);
    }
    
    TValue& Add(const TKey &Key, const TValue &Value, uint32 BucketIndex) {
        Check(BucketIndex == ComputeBucket(Key));
        Check(!Contains(Key, BucketIndex));
        if (IsFull()) {
            Grow();
            BucketIndex = ComputeBucket(Key);
        }
        return AddImpl(Key, Value, BucketIndex);
    }

    TValue& Add(const TKey &Key, TValue &&Value) {
        Check(!Contains(Key));
        if (IsFull()) {
            Grow();
        }
        auto BucketIndex = ComputeBucket(Key);
        return AddImpl(Key, Move(Value), BucketIndex);
    }

    TValue& Add(const TKey &Key, TValue &&Value, uint32 BucketIndex) {
        Check(BucketIndex == ComputeBucket(Key));
        Check(!Contains(Key, BucketIndex));
        if (IsFull()) {
            Grow();
            BucketIndex = ComputeBucket(Key);
        }
        return AddImpl(Key, Move(Value), BucketIndex);
    }
    
    TValue& Set(const TKey &Key, TValue &&Value) {
        if (IsFull()) {
            Grow();
        }
        auto BucketIndex = ComputeBucket(Key);
        for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
            if (*Entry->Key == Key) {
                *Entry->Value = Move(Value);
                return *Entry->Value;
            }
        }
        return AddImpl(Key, Move(Value), BucketIndex);
    }
    
    TValue& Set(const TKey &Key, TValue &&Value, uint32 BucketIndex) {
        Check(BucketIndex == ComputeBucket(Key));
        for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
            if (*Entry->Key == Key) {
                *Entry->Value = Move(Value);
                return *Entry->Value;
            }
        }
        if (IsFull()) {
            Grow();
            BucketIndex = ComputeBucket(Key);
        }
        return AddImpl(Key, Move(Value), BucketIndex);
    }
    
    TValue& Set(const TKey &Key, const TValue &Value) {
        if (IsFull()) {
            Grow();
        }
        auto BucketIndex = ComputeBucket(Key);
        for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
            if (*Entry->Key == Key) {
                *Entry->Value = Value;
                return *Entry->Value;
            }
        }
        return AddImpl(Key, Value, BucketIndex);
    }
    
    TValue& Set(const TKey &Key, const TValue &Value, uint32 BucketIndex) {
        Check(BucketIndex == ComputeBucket(Key));
        for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
            if (*Entry->Key == Key) {
                *Entry->Value = Value;
                return *Entry->Value;
            }
        }
        if (IsFull()) {
            Grow();
            BucketIndex = ComputeBucket(Key);
        }
        return AddImpl(Key, Value, BucketIndex);
    }
    
    TValue &GetOrCreate(const TKey &Key) {
        auto BucketIndex = ComputeBucket(Key);
        if (Entries.GetSize() > BucketIndex + 1) {
            for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
                if (*Entry->Key == Key) {
                    return *Entry->Value;
                }
            }
        }
        return Create(Key, BucketIndex);
    }

    TValue &GetOrCreate(const TKey &Key, uint32 BucketIndex) {
        Check(BucketIndex == ComputeBucket(Key));
        if (Entries.GetSize() > BucketIndex + 1) {
            for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
                if (*Entry->Key == Key) {
                    return *Entry->Value;
                }
            }
        }
        return Create(Key, BucketIndex);
    }
    
    TValue &GetOrAdd(const TKey &Key, TValue &&DefaultValue) {
        auto BucketIndex = ComputeBucket(Key);
        if (Entries.GetSize() > BucketIndex + 1) {
            for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
                if (*Entry->Key == Key) {
                    return *Entry->Value;
                }
            }
        }
        return Add(Key, Move(DefaultValue), BucketIndex);
    }
    
    TValue &GetOrAdd(const TKey &Key, TValue &&DefaultValue, uint32 BucketIndex) {
        Check(BucketIndex == ComputeBucket(Key));
        if (Entries.GetSize() > BucketIndex + 1) {
            for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
                if (*Entry->Key == Key) {
                    return *Entry->Value;
                }
            }
        }
        return Add(Key, Move(DefaultValue), BucketIndex);
    }
    
    TValue &GetOrAdd(const TKey &Key, const TValue &DefaultValue) {
        auto BucketIndex = ComputeBucket(Key);
        if (Entries.GetSize() > BucketIndex + 1) {
            for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
                if (*Entry->Key == Key) {
                    return *Entry->Value;
                }
            }
        }
        return Add(Key, DefaultValue, BucketIndex);
    }
    
    TValue &GetOrAdd(const TKey &Key, const TValue &DefaultValue, uint32 BucketIndex) {
        Check(BucketIndex == ComputeBucket(Key));
        if (Entries.GetSize() > BucketIndex + 1) {
            for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
                if (*Entry->Key == Key) {
                    return *Entry->Value;
                }
            }
        }
        return Add(Key, DefaultValue, BucketIndex);
    }

    TValue &Get(const TKey &Key) {
        Check(Contains(Key));
        auto BucketIndex = ComputeBucket(Key);
        for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
            if (*Entry->Key == Key) {
                return *Entry->Value;
            }
        }
    }

    TValue &Get(const TKey &Key, uint32 BucketIndex) {
        Check(ComputeBucket(Key) == BucketIndex);
        Check(Contains(Key, BucketIndex));
        for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
            if (*Entry->Key == Key) {
                return *Entry->Value;
            }
        }
    }

    const TValue &Get(const TKey &Key) const {
        return ((THashMap *)this)->Get(Key);
    }

    const TValue &Get(const TKey &Key, uint32 BucketIndex) const {
        return ((THashMap *)this)->Get(Key, BucketIndex);
    }

    TValue &operator[](const TKey &Key) {
        auto BucketIndex = ComputeBucket(Key);
        return *FindEntry(Key, BucketIndex)->Value;
    }

    const TValue &operator[](const TKey &Key) const {
        return ((THashMap *)this)->operator[](Key);
    }

    TValue *Find(const TKey &Key) {
        auto BucketIndex = ComputeBucket(Key);
        FEntry *Entry = FindEntry(Key, BucketIndex);
        if (Entry) {
            return &*Entry->Value;
        }
        return nullptr;
    }

    TValue *Find(const TKey &Key, uint32 BucketIndex) {
        Check(ComputeBucket(Key) == BucketIndex);
        FEntry *Entry = FindEntry(Key, BucketIndex);
        if (Entry) {
            return &*Entry->Value;
        }
        return nullptr;
    }

    const TValue *Find(const TKey &Key) const {
        return ((THashMap *)this)->Find(Key);
    }

    const TValue *Find(const TKey &Key, uint32 BucketIndex) const {
        return ((THashMap *)this)->Find(Key, BucketIndex);
    }
    
    void Remove(const TKey &Key) {
        Check(Contains(Key));
        auto BucketIndex = ComputeBucket(Key);
        FEntry **Bucket = &Entries[BucketIndex];
        for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
            if (*Entry->Key == Key) {
                FEntry **Previous = Bucket;
                while (*Previous != Entry)
                    Previous = &(*Previous)->Next;
                *Previous = Entry->Next;
                if constexpr (!Meta::IsTriviallyDestructible<TKey>) {
                    Meta::Destruct(*Entry->Key);
                }
                if constexpr (!Meta::IsTriviallyDestructible<TValue>) {
                    Meta::Destruct(*Entry->Value);
                }
                Allocator.Free(Entry);
                --ElementCount;
                return;
            }
        }
    }
    
    void Remove(const TKey &Key, uint32 BucketIndex) {
        Check(ComputeBucket(Key) == BucketIndex);
        Check(Contains(Key, BucketIndex));
        FEntry **Bucket = &Entries[BucketIndex];
        for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
            if (*Entry->Key == Key) {
                FEntry **Previous = Bucket;
                while (*Previous != Entry)
                    Previous = &(*Previous)->Next;
                *Previous = Entry->Next;
                if constexpr (!Meta::IsTriviallyDestructible<TKey>) {
                    Meta::Destruct(*Entry->Key);
                }
                if constexpr (!Meta::IsTriviallyDestructible<TValue>) {
                    Meta::Destruct(*Entry->Value);
                }
                Allocator.Free(Entry);
                --ElementCount;
                return;
            }
        }
    }

    void Clear() {
        FEntry **End = &Entries.Last();
        for (FEntry **Bucket = &Entries.First(); Bucket != End; ++Bucket) {
            for (FEntry *Entry = *Bucket; Entry;) {
                FEntry *Next = Entry->Next;
                if constexpr (!Meta::IsTriviallyDestructible<TKey>) {
                    Meta::Destruct(*Entry->Key);
                }
                if constexpr (!Meta::IsTriviallyDestructible<TValue>) {
                    Meta::Destruct(*Entry->Value);
                }
                Allocator.Free(Entry);
                Entry = Next;
            }
            *Bucket = nullptr;
        }
        ElementCount = 0;
    }
    
    bool IsEmpty() const {
        return ElementCount == 0;
    }

private:
    // Last entry is an end iterator dummy (value of 1).
    TArray<FEntry *, FHeapAllocator> Entries;
    uint32 ElementCount = 0;
    int8 BucketHashShift = 63;
    FAllocator Allocator;

    TValue &CreateImpl(const TKey &Key, const uint32 BucketIndex) {
        FEntry *&Bucket = Entries[BucketIndex];
        FEntry *Entry = (FEntry *)Allocator.Alloc();
        if constexpr (Meta::IsTrivialCopy<TKey>) {
            *Entry->Key = Key;
        } else {
            Meta::Construct(*Entry->Key, Key);
        }
        if constexpr (!Meta::IsTriviallyDefaultConstructible<TValue>) {
            Meta::Construct(*Entry->Value);
        }
        Entry->Next = Bucket;
        Bucket = Entry;
        ++ElementCount;
        return *Entry->Value;
    }

    TValue& AddImpl(const TKey &Key, const TValue &Value, uint32 BucketIndex) {
        FEntry *&Bucket = Entries[BucketIndex];
        FEntry *Entry = (FEntry *)Allocator.Alloc();
        if constexpr (Meta::IsTrivialCopy<TKey>) {
            *Entry->Key = Key;
        } else {
            Meta::Construct(*Entry->Key, Key);
        }
        if constexpr (Meta::IsTrivialCopy<TValue>) {
            *Entry->Value = Value;
        } else {
            Meta::Construct(*Entry->Value, Value);
        }
        Entry->Next = Bucket;
        Bucket = Entry;
        ++ElementCount;
        return *Entry->Value;
    }
    
    TValue& AddImpl(const TKey &Key, TValue &&Value, uint32 BucketIndex) {
        FEntry *&Bucket = Entries[BucketIndex];
        FEntry *Entry = (FEntry *)Allocator.Alloc();

        if constexpr (Meta::IsTrivialCopy<TKey>) {
            *Entry->Key = Key;
        } else {
            Meta::Construct(*Entry->Key, Key);
        }
        if constexpr (Meta::IsTrivialMove<TValue>) {
            *Entry->Value = Move(Value);
        } else {
            Meta::Construct(*Entry->Value, Move(Value));
        }
        Entry->Next = Bucket;
        Bucket = Entry;
        ++ElementCount;
        return *Entry->Value;
    }

    void Construct(const THashMap &Other) {
        Entries.Resize(Other.Entries.GetSize(), nullptr);
        Entries.Last() = (FEntry *)1;
        for (uint32 Index = 0, End = Other.Entries.GetSize() - 1; Index < End; ++Index) {
            for (FEntry *Entry = (FEntry *)Other.Entries[Index]; Entry; Entry = Entry->Next) {
                FEntry *NewEntry = (FEntry *)Allocator.Alloc();
                if constexpr (Meta::IsTrivialCopy<TKey>) {
                    *NewEntry->Key = *Entry->Key;
                } else {
                    Meta::Construct(*NewEntry->Key, *Entry->Key);
                }
                if constexpr (Meta::IsTrivialCopy<TValue>) {
                    *NewEntry->Value = *Entry->Value;
                } else {
                    Meta::Construct(*NewEntry->Value, *Entry->Value);
                }
                NewEntry->Next = Entries[Index];
                Entries[Index] = NewEntry;
            }
        }
        ElementCount = Other.ElementCount;
        BucketHashShift = Other.BucketHashShift;
    }

    void MoveConstruct(THashMap &&Other) {
        Swap(Entries, Other.Entries);
        Swap(ElementCount, Other.ElementCount);
        Swap(BucketHashShift, Other.BucketHashShift);
        Swap(Allocator, Other.Allocator);
    }

    FEntry *FindEntry(const TKey &Key, uint32 BucketIndex) {
        if (Entries.GetSize() > BucketIndex + 1) {
            for (FEntry *Entry = Entries[BucketIndex]; Entry; Entry = Entry->Next) {
                if (*Entry->Key == Key) {
                    return Entry;
                }
            }
        }
        return nullptr;
    }

    const FEntry *FindEntry(const TKey &Key, uint32 BucketIndex) const {
        return ((THashMap *)this)->FindEntry(Key, BucketIndex);
    }

    uint32 BucketFromHash(hash InHash) const {
        return (11400714819323198485ull * InHash) >> BucketHashShift;
    }

    bool IsFull() const {
        return ElementCount + 1 >= Entries.GetSize();
    }

    void Rehash(uint32 InBucketCount) {
        uint32 BucketCount = Math::Max(InBucketCount, ElementCount);
        Check(BucketCount > 0);
        if (BucketCount == 0) {
            Entries.Clear();
            BucketHashShift = 63;
            return;
        }
        BucketCount = Math::Max((uint32)2, Math::NextPowerOfTwo(InBucketCount));
        if (BucketCount + 1 == Entries.GetSize()) {
            return;
        }
        BucketHashShift = INT8_C(64) - Math::Log2(BucketCount);
        TArray<FEntry *, FHeapAllocator> NewEntries;
        NewEntries.Resize(BucketCount + 1, nullptr);
        NewEntries.Last() = (FEntry *)1;
        Swap(Entries, NewEntries);
        // Remap old bucket entries to new buckets
        for (uint32 Index = 0; Index < NewEntries.GetSize() - 1; ++Index) {
            for (FEntry *Entry = NewEntries[Index]; Entry;) {
                FEntry *Next = Entry->Next;
                const uint32 BucketIndex = BucketFromHash(THasher::Get(*Entry->Key));
                FEntry **Bucket = &Entries[BucketIndex];
                Entry->Next = *Bucket;
                *Bucket = Entry;
                Entry = Next;
            }
        }
    }

    void Grow() {
        Check(IsFull());
        Rehash(Math::Max((uint32)4, Entries.GetSize() * 2));
    }
};

}  // namespace SE

HashMapIterators.h

#pragma once

#include "Array/ArrayUtility.h"
#include "Hash/HashUtility.h"
#include "Math/Math.h"

namespace SE {

namespace Impl {
template<typename TKey, typename TValue>
struct THashMapEntry {
    THashMapEntry *Next;
    TUninitialized<TKey> Key;
    TUninitialized<TValue> Value;
};
}  // namespace Impl

template<typename TKey, typename TValue>
class THashMapConstIterator {
private:
    using FEntry = Impl::THashMapEntry<TKey, TValue>;

public:
    THashMapConstIterator() {
    }

    THashMapConstIterator(FEntry *CurrentEntry, FEntry *const *CurrentBucket) : CurrentEntry(CurrentEntry), CurrentBucket(CurrentBucket) {
    }

    friend bool operator==(const THashMapConstIterator &Left, const THashMapConstIterator &Right) {
        return Left.CurrentEntry == Right.CurrentEntry;
    }

    friend bool operator!=(const THashMapConstIterator &Left, const THashMapConstIterator &Right) {
        return !(Left == Right);
    }

    const TKey &GetKey() const {
        Check(CurrentEntry > (FEntry *)1);
        return *CurrentEntry->Key;
    }

    const TValue &GetValue() const {
        Check(CurrentEntry > (FEntry *)1);
        return *CurrentEntry->Value;
    }

    THashMapConstIterator &operator++() {
        Check(CurrentEntry);
        if (CurrentEntry->Next) {
            CurrentEntry = CurrentEntry->Next;
        } else {
            do {
                ++CurrentBucket;
            } while (!*CurrentBucket);
            CurrentEntry = *CurrentBucket;
        }
        return *this;
    }

    THashMapConstIterator operator++(int) {
        THashMapConstIterator Copy(*this);
        ++*this;
        return Copy;
    }

    operator bool() const {
        return CurrentEntry != (FEntry *)1;
    }

private:
    FEntry *CurrentEntry = 1;
    FEntry *const *CurrentBucket = 1;
};

template<typename TKey, typename TValue, typename TAllocator>
class THashMapIterator {
private:
    using FEntry = Impl::THashMapEntry<TKey, TValue>;

public:
    THashMapIterator() {
    }

    THashMapIterator(uint32 *ElementCount, FEntry *CurrentEntry, FEntry **CurrentBucket, TAllocator *Allocator)
        : ElementCount(ElementCount),
          CurrentEntry(CurrentEntry),
          CurrentBucket(CurrentBucket),
          Allocator(Allocator) {
    }

    friend bool operator==(const THashMapIterator &Left, const THashMapIterator &Right) {
        return Left.CurrentEntry == Right.CurrentEntry;
    }

    friend bool operator!=(const THashMapIterator &Left, const THashMapIterator &Right) {
        return !(Left == Right);
    }

    operator THashMapConstIterator<TKey, TValue>() const {
        return {CurrentEntry, CurrentBucket};
    }

    const TKey &GetKey() const {
        Check(CurrentEntry > (FEntry *)1);
        Check(CurrentEntry != &IntermediateEntry);
        return *CurrentEntry->Key;
    }

    TValue &GetValue() const {
        Check(CurrentEntry > (FEntry *)1);
        Check(CurrentEntry != &IntermediateEntry);
        return *CurrentEntry->Value;
    }

    THashMapIterator &operator++() {
        Check(CurrentEntry);
        if (CurrentEntry->Next) {
            CurrentEntry = CurrentEntry->Next;
        } else {
            do {
                ++CurrentBucket;
            } while (!*CurrentBucket);
            CurrentEntry = *CurrentBucket;
        }
        return *this;
    }

    THashMapIterator operator++(int) {
        THashMapIterator Copy(*this);
        ++*this;
        return Copy;
    }

    operator bool() const {
        return CurrentEntry != (FEntry *)1;
    }

    // Increment the iterator after this function before accessing it.
    void RemoveCurrent() {
        Check(ElementCount);
        Check(CurrentEntry);
        Check(CurrentEntry != &IntermediateEntry);
        // Either entry array element's bucket pointer, or entry's linked list next pointer.
        FEntry **Previous = CurrentBucket;
        FEntry *Entry = CurrentEntry;
        while (*Previous != Entry) {
            Previous = &(*Previous)->Next;
        }
        *Previous = Entry->Next;
        if (Entry->Next) {
            IntermediateEntry.Next = Entry->Next;
        } else {
            do {
                ++CurrentBucket;
            } while (!*CurrentBucket);
            IntermediateEntry.Next = *CurrentBucket;
        }
        CurrentEntry = &IntermediateEntry;
        if constexpr (!Meta::IsTriviallyDestructible<TKey>) {
            Meta::Destruct(*Entry->Key);
        }
        if constexpr (!Meta::IsTriviallyDestructible<TValue>) {
            Meta::Destruct(*Entry->Value);
        }
        Allocator->Free(Entry);
        --(*ElementCount);
    }

private:
    uint32 *ElementCount = nullptr;
    FEntry *CurrentEntry = 1;
    FEntry **CurrentBucket = 1;
    TAllocator *Allocator = nullptr;
    // Used as intermediate entry active after RemoveCurrent call.
    FEntry IntermediateEntry;
};

}  // namespace SE