StripedMap.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "CoreTypes.h"
 
#include "Containers/Map.h"
#include "Async/SharedLock.h"
#include "Async/SharedMutex.h"
#include "Async/UniqueLock.h"
#include "HAL/PlatformTLS.h"
#include "Misc/TransactionallySafeRWLock.h"
#include "Misc/ScopeRWLock.h"
 
struct FSharedMutexStripedMapLockingPolicy
{
    typedef UE::FSharedMutex                  MutexType;
    typedef UE::TUniqueLock<UE::FSharedMutex> ExclusiveLockType;
    typedef UE::TSharedLock<UE::FSharedMutex> SharedLockType;
};
 
struct FTransactionallySafeStripedMapLockingPolicy
{
    typedef FTransactionallySafeRWLock        MutexType;
    typedef UE::TWriteScopeLock<MutexType>    ExclusiveLockType;
    typedef UE::TReadScopeLock<MutexType>     SharedLockType;
};
 
// Use FTransactionallySafeRWLock by default for now because there is no transactionally-safe
// FSharedMutex and a bug in FSharedMutex is currently causing deadlocks.
typedef FTransactionallySafeStripedMapLockingPolicy FDefaultStripedMapLockingPolicy;
 
template <int32 BucketCount, typename BaseMapType, typename KeyType, typename ValueType, typename SetAllocator, typename KeyFuncs, typename LockingPolicy>
class TStripedMapBase
{
    static_assert(BucketCount > 0, "The BucketCount needs to be at least 1");
public:
    typedef BaseMapType                               MapType;
    typedef typename BaseMapType::KeyConstPointerType KeyConstPointerType;
    typedef KeyFuncs                                  KeyFuncsType;
 
protected:
    struct FDebuggableMutex
    {
        typename LockingPolicy::MutexType Mutex;
        uint32 ExclusiveLockOwnerThreadId   = 0;
        std::atomic<uint32> SharedLockCount = 0;
    };
 
    struct FDebuggableSharedLock
    {
        FDebuggableMutex& DebuggableMutex;
        typename LockingPolicy::SharedLockType InnerLock;
        FDebuggableSharedLock(FDebuggableMutex& InDebuggableMutex)
            : DebuggableMutex(InDebuggableMutex)
            , InnerLock(InDebuggableMutex.Mutex)
        {
            check(DebuggableMutex.ExclusiveLockOwnerThreadId == 0);
            DebuggableMutex.SharedLockCount++;
        }
 
        ~FDebuggableSharedLock()
        {
            DebuggableMutex.SharedLockCount--;
        }
    };
 
    struct FDebuggableExclusiveLock
    {
        FDebuggableMutex& DebuggableMutex;
        typename LockingPolicy::ExclusiveLockType InnerLock;
        FDebuggableExclusiveLock(FDebuggableMutex& InDebuggableMutex)
            : DebuggableMutex(InDebuggableMutex)
            , InnerLock(InDebuggableMutex.Mutex)
        {
            check(DebuggableMutex.ExclusiveLockOwnerThreadId == 0);
            DebuggableMutex.ExclusiveLockOwnerThreadId = FPlatformTLS::GetCurrentThreadId();
        }
 
        ~FDebuggableExclusiveLock()
        {
            DebuggableMutex.ExclusiveLockOwnerThreadId = 0;
        }
    };
 
#define WITH_STRIPEDMAP_DEBUGGABLE_MUTEX 0
#if WITH_STRIPEDMAP_DEBUGGABLE_MUTEX
    typedef FDebuggableMutex                          MutexType;
    typedef FDebuggableExclusiveLock                  ExclusiveLockType;
    typedef FDebuggableSharedLock                     SharedLockType;
#else
    typedef typename LockingPolicy::MutexType         MutexType;
    typedef typename LockingPolicy::ExclusiveLockType ExclusiveLockType;
    typedef typename LockingPolicy::SharedLockType    SharedLockType;
#endif
#undef WITH_STRIPEDMAP_DEBUGGABLE_MUTEX
 
    struct FBucket
    {
        mutable MutexType Lock;
        MapType           Map;
    } Buckets[BucketCount];
 
    uint32 GetBucketIndex(uint32 Hash) const
    {
        if constexpr (BucketCount == 1)
        {
            return 0;
        }
        else
        {
            return Hash % BucketCount;
        }
    }
 
    template <typename FunctionType>
    decltype(auto) ApplyUnlockedByHash(uint32 InHash, FunctionType&& InFunction)
    {
        return InFunction(InHash, Buckets[GetBucketIndex(InHash)]);
    }
 
    template <typename FunctionType>
    decltype(auto) ApplyUnlocked(KeyConstPointerType InKey, FunctionType&& InFunction)
    {
        return ApplyUnlockedByHash(KeyFuncsType::GetKeyHash(InKey), Forward<FunctionType>(InFunction));
    }
 
    template <typename LockType, typename FunctionType>
    decltype(auto) ApplyByHash(uint32 InHash, FunctionType&& InFunction) const
    {
        const FBucket& Bucket = Buckets[GetBucketIndex(InHash)];
 
        LockType ScopeLock(Bucket.Lock);
        return InFunction(InHash, Bucket.Map);
    }
 
    template <typename LockType, typename FunctionType>
    decltype(auto) Apply(KeyConstPointerType InKey, FunctionType&& InFunction) const
    {
        return ApplyByHash<LockType>(KeyFuncsType::GetKeyHash(InKey), Forward<FunctionType>(InFunction));
    }
 
    template <typename LockType, typename FunctionType>
    decltype(auto) ApplyByHash(uint32 InHash, FunctionType&& InFunction)
    {
        FBucket& Bucket = Buckets[GetBucketIndex(InHash)];
 
        LockType ScopeLock(Bucket.Lock);
        return InFunction(InHash, Bucket.Map);
    }
 
    template <typename LockType, typename FunctionType>
    decltype(auto) Apply(KeyConstPointerType InKey, FunctionType&& InFunction)
    {
        return ApplyByHash<LockType>(KeyFuncsType::GetKeyHash(InKey), Forward<FunctionType>(InFunction));
    }
 
    template <typename FunctionType>
    decltype(auto) Write(KeyConstPointerType InKey, FunctionType&& InFunction)
    {
        return Apply<ExclusiveLockType>(InKey, Forward<FunctionType>(InFunction));
    }
 
    template <typename FunctionType>
    decltype(auto) WriteByHash(uint32 InHash, KeyConstPointerType InKey, FunctionType&& InFunction)
    {
        return ApplyByHash<ExclusiveLockType>(InHash, InKey, Forward<FunctionType>(InFunction));
    }
 
    template <typename FunctionType>
    decltype(auto) Read(KeyConstPointerType InKey, FunctionType&& InFunction)
    {
        return Apply<SharedLockType>(InKey, Forward<FunctionType>(InFunction));
    }
 
    template <typename FunctionType>
    decltype(auto) ReadByHash(uint32 InHash, KeyConstPointerType InKey, FunctionType&& InFunction)
    {
        return ApplyByHash<SharedLockType>(InHash, InKey, Forward<FunctionType>(InFunction));
    }
 
    template <typename FunctionType>
    decltype(auto) Read(KeyConstPointerType InKey, FunctionType&& InFunction) const
    {
        return Apply<SharedLockType>(InKey, Forward<FunctionType>(InFunction));
    }
 
    template <typename FunctionType>
    decltype(auto) ReadByHash(uint32 InHash, KeyConstPointerType InKey, FunctionType&& InFunction) const
    {
        return ApplyByHash<SharedLockType>(InHash, InKey, Forward<FunctionType>(InFunction));
    }
 
    template <typename FunctionType>
    void ForEachMap(FunctionType&& InFunction)
    {
        for (FBucket& Bucket : Buckets)
        {
            ExclusiveLockType ScopeLock(Bucket.Lock);
            InFunction(Bucket.Map);
        }
    }
 
    template <typename FunctionType>
    void ForEachMap(FunctionType&& InFunction) const
    {
        for (const FBucket& Bucket : Buckets)
        {
            SharedLockType ScopeLock(Bucket.Lock);
            InFunction(Bucket.Map);
        }
    }
 
public:
    ValueType FindRef(KeyConstPointerType InKey) const
    {
        return Read(InKey, [&InKey](uint32 Hash, const MapType& Map)
            {
                const ValueType* Result = Map.FindByHash(Hash, InKey);
                if (Result)
                {
                    return *Result;
                }
                
                return ValueType();
            }
        );
    }
 
    template <typename FunctionType>
    bool FindAndApply(KeyConstPointerType InKey, FunctionType&& InFunction) const
    {
        return Read(InKey, [&InKey, &InFunction](uint32 Hash, const MapType& Map)
            {
                const ValueType* Result = Map.FindByHash(Hash, InKey);
                if (Result)
                {
                    InFunction(*Result);
                    return true;
                }
 
                return false;
            }
        );
    }
 
    template <typename UpdateFunctionType>
    bool FindAndApply(KeyConstPointerType InKey, UpdateFunctionType&& InUpdateFunction)
    {
        return Write(InKey, [&InKey, &InUpdateFunction](uint32 Hash, MapType& Map)
            {
                ValueType* Result = Map.FindByHash(Hash, InKey);
                if (Result)
                {
                    InUpdateFunction(*Result);
                    return true;
                }
 
                return false;
            }
        );
    }
 
    [[nodiscard]] bool Contains(KeyConstPointerType InKey) const
    {
        return Read(InKey, [&InKey](uint32 Hash, const MapType& Map) { return Map.ContainsByHash(Hash, InKey); });
    }
 
    template <typename InitKeyType = KeyType, typename InitValueType = ValueType>
    void Emplace(InitKeyType&& InKey, InitValueType&& InValue)
    {
        return Write(InKey, [&InKey, &InValue](uint32 Hash, MapType& Map) 
            {
                Map.EmplaceByHash(Hash, Forward<InitKeyType>(InKey), Forward<InitValueType>(InValue));
            }
        );
    }
 
    void Add(const KeyType& InKey, const ValueType& InValue)
    {
        return Emplace(InKey, InValue);
    }                                                      
                
    void Add(const KeyType& InKey, ValueType&& InValue)
    {                                                      
        return Emplace(InKey, MoveTempIfPossible(InValue));
    }                                                      
    
    void Add(KeyType&& InKey, const ValueType& InValue)
    {                                                      
        return Emplace(MoveTempIfPossible(InKey), InValue);
    }
 
    void Add(KeyType&& InKey, ValueType&& InValue)
    {                                                      
        return Emplace(MoveTempIfPossible(InKey), MoveTempIfPossible(InValue));
    }
 
    template <typename ProduceFunctionType>
    ValueType FindOrProduce(const KeyType& InKey, ProduceFunctionType&& InProduceFunction)
    {
        return ApplyUnlocked(InKey,
            [this, &InKey, &InProduceFunction](uint32 Hash, FBucket& Bucket)
            {
                {
                    SharedLockType ScopeLock(Bucket.Lock);
                    const ValueType* Result = Bucket.Map.FindByHash(Hash, InKey);
                    if (Result)
                    {
                        return *Result;
                    }
                }
 
                ExclusiveLockType ScopeLock(Bucket.Lock);
                ValueType* Result = Bucket.Map.FindByHash(Hash, InKey);
                if (Result)
                {
                    return *Result;
                }
 
                return Bucket.Map.AddByHash(Hash, InKey, InProduceFunction());
            });
    }
 
    template <typename ProduceFunctionType, typename ApplyFunctionType>
    void FindOrProduceAndApply(const KeyType& InKey, ProduceFunctionType&& InProduceFunction, ApplyFunctionType&& InApplyFunction)
    {
        ApplyUnlocked(InKey,
            [this, &InKey, &InApplyFunction, &InProduceFunction](uint32 Hash, FBucket& Bucket)
            {
                {
                    SharedLockType ScopeLock(Bucket.Lock);
                    const ValueType* Result = Bucket.Map.FindByHash(Hash, InKey);
                    if (Result)
                    {
                        InApplyFunction(*Result);
                        return;
                    }
                }
 
                ExclusiveLockType ScopeLock(Bucket.Lock);
                ValueType* Result = Bucket.Map.FindByHash(Hash, InKey);
                if (Result)
                {
                    InApplyFunction(*Result);
                    return;
                }
 
                InApplyFunction(Bucket.Map.AddByHash(Hash, InKey, InProduceFunction()));
            });
    }
 
    template <typename TryProduceFunctionType, typename ApplyFunctionType>
    bool FindOrTryProduceAndApply(const KeyType& InKey, TryProduceFunctionType&& InTryProduceFunction, ApplyFunctionType&& InApplyFunction)
    {
        return ApplyUnlocked(InKey,
            [this, &InKey, &InApplyFunction, &InTryProduceFunction](uint32 Hash, FBucket& Bucket)
            {
                {
                    SharedLockType ScopeLock(Bucket.Lock);
                    const ValueType* Result = Bucket.Map.FindByHash(Hash, InKey);
                    if (Result)
                    {
                        InApplyFunction(*Result);
                        return true;
                    }
                }
 
                ExclusiveLockType ScopeLock(Bucket.Lock);
                const ValueType* Result = Bucket.Map.FindByHash(Hash, InKey);
                if (Result)
                {
                    InApplyFunction(*Result);
                    return true;
                }
 
                ValueType Value;
                if (InTryProduceFunction(Value))
                {
                    InApplyFunction(Bucket.Map.AddByHash(Hash, InKey, MoveTemp(Value)));
                    return true;
                }
 
                return false;
            });
    }
 
    template <typename ProduceFunctionType, typename ApplyFunctionType>
    void FindOrProduceAndApplyForWrite(const KeyType& InKey, ProduceFunctionType&& InProduceFunction, ApplyFunctionType&& InApplyFunction)
    {
        return Write(InKey, [&InKey, &InProduceFunction, &InApplyFunction](uint32 Hash, MapType& Map)
        {
            ValueType* Result = Map.FindByHash(Hash, InKey);
            if (Result)
            {
                InApplyFunction(*Result);
            }
            else
            {
                InApplyFunction(Map.AddByHash(Hash, InKey, InProduceFunction()));
            }
        });
    }
 
    template <typename TryProduceFunctionType, typename ApplyFunctionType>
    bool FindOrTryProduceAndApplyForWrite(const KeyType& InKey, TryProduceFunctionType&& InTryProduceFunction, ApplyFunctionType&& InApplyFunction)
    {
        return Write(InKey, [&InKey, &InTryProduceFunction, &InApplyFunction](uint32 Hash, MapType& Map)
        {
            ValueType* Result = Map.FindByHash(Hash, InKey);
            if (Result)
            {
                InApplyFunction(*Result);
                return true;
            }
 
            ValueType Value;
            if (InTryProduceFunction(Value))
            {
                InApplyFunction(Map.AddByHash(Hash, InKey, MoveTemp(Value)));
                return true;
            }
 
            return false;
        });
    }
 
    int32 Remove(KeyConstPointerType InKey)
    {
        return Write(InKey, [&InKey](uint32 Hash, MapType& Map) { return Map.RemoveByHash(Hash, InKey); });
    }
    
    int32 RemoveByHash(uint32 InKeyHash, KeyConstPointerType InKey)
    {
        return WriteByHash(InKeyHash, InKey, [&InKey](uint32 Hash, MapType& Map) { return Map.RemoveByHash(Hash, InKey); });
    }
 
    template <typename PredicateType>
    int32 RemoveIf(KeyConstPointerType InKey, PredicateType&& InPredicate)
    {
        return Write(InKey,
            [&InKey, &InPredicate](uint32 Hash, MapType& Map)
            {
                FSetElementId ElementId = Map.FindIdByHash(Hash, InKey);
                if (ElementId.IsValidId() && InPredicate(Map.Get(ElementId).Value))
                {
                    Map.Remove(ElementId);
                    return 1;
                }
 
                return 0;
            });
    }
 
    template <typename PredicateType>
    int32 RemoveIf(PredicateType&& InPredicate)
    {
        int32 RemovedCount = 0;
        ForEachMap(
            [&RemovedCount, &InPredicate](MapType& Map)
            {
                for (auto It = Map.CreateIterator(); It; ++It)
                {
                    if (InPredicate(*It))
                    {
                        It.RemoveCurrent();
                        RemovedCount++;
                    }
                }
            }
        );
        return RemovedCount;
    }
 
    bool RemoveAndCopyValue(KeyConstPointerType InKey, ValueType& OutRemovedValue)
    {
        return Write(InKey, [&InKey, &OutRemovedValue](uint32 Hash, MapType& Map) 
            {
                return Map.RemoveAndCopyValueByHash(Hash, InKey, OutRemovedValue); 
            });
    }
 
    ValueType FindAndRemoveChecked(KeyConstPointerType InKey)
    {
        return Write(InKey, [&InKey](uint32 Hash, MapType& Map)
            {
                return Map.FindAndRemoveChecked(InKey); 
            }
        );
    }
 
    void Empty()
    {
        ForEachMap(
            [](MapType& Map)
            {
                Map.Empty();
            }
        );
    }
 
    void Reset()
    {
        ForEachMap(
            [](MapType& Map)
            {
                Map.Reset();
            }
        );
    }
 
    void Shrink()
    {
        ForEachMap(
            [](MapType& Map)
            {
                Map.Shrink();
            }
        );
    }
 
    void Compact()
    {
        ForEachMap(
            [](MapType& Map)
            {
                Map.Compact();
            }
        );
    }
 
    [[nodiscard]] int32 Num() const
    {
        int32 Count = 0;
        ForEachMap(
            [&Count](const MapType& Map)
            {
                Count += Map.Num();
            }
        );
 
        return Count;
    }
 
    template <typename FunctionType>
    void ForEach(FunctionType&& InFunction)
    {
        ForEachMap(
            [&InFunction](MapType& Map)
            {
                for (auto& Item : Map)
                {
                    InFunction(Item);
                }
            }
        );
    }
 
    template <typename FunctionType>
    void ForEach(FunctionType&& InFunction) const
    {
        ForEachMap(
            [&InFunction](const MapType& Map)
            {
                for (const auto& Item : Map)
                {
                    InFunction(Item);
                }
            }
        );
    }
 
    template<typename Allocator>
    int32 GetKeys(TArray<KeyType, Allocator>& OutKeys) const
    {
        OutKeys.Reset();
        TSet<KeyType> VisitedKeys;
        ForEach(
            [&OutKeys, &VisitedKeys](const auto& Pair)
            {
                if (!VisitedKeys.Contains(Pair.Key))
                {
                    OutKeys.Add(Pair.Key);
                    VisitedKeys.Add(Pair.Key);
                }
            }
        );
        return OutKeys.Num();
    }
};
 
template <
    int32 BucketCount, 
    typename KeyType, 
    typename ValueType, 
    typename SetAllocator = FDefaultSetAllocator, 
    typename KeyFuncs = TDefaultMapHashableKeyFuncs<KeyType, ValueType, false>, 
    typename LockingPolicy = FDefaultStripedMapLockingPolicy
>
class TStripedMap : public TStripedMapBase<BucketCount, TMap<KeyType, ValueType, SetAllocator, KeyFuncs>, KeyType, ValueType, SetAllocator, KeyFuncs, LockingPolicy>
{
public:
    using Super = TStripedMapBase<BucketCount, TMap<KeyType, ValueType, SetAllocator, KeyFuncs>, KeyType, ValueType, SetAllocator, KeyFuncs, LockingPolicy>;
};
 
template <
    int32 BucketCount, 
    typename KeyType, 
    typename ValueType, 
    typename SetAllocator = FDefaultSetAllocator, 
    typename KeyFuncs = TDefaultMapHashableKeyFuncs<KeyType, ValueType, true>, 
    typename LockingPolicy = FDefaultStripedMapLockingPolicy
>
class TStripedMultiMap : public TStripedMapBase<BucketCount, TMultiMap<KeyType, ValueType, SetAllocator, KeyFuncs>, KeyType, ValueType, SetAllocator, KeyFuncs, LockingPolicy>
{
public:
    using Super = TStripedMapBase<BucketCount, TMultiMap<KeyType, ValueType, SetAllocator, KeyFuncs>, KeyType, ValueType, SetAllocator, KeyFuncs, LockingPolicy>;
 
    template<typename Allocator>
    void MultiFind(typename Super::KeyConstPointerType InKey, TArray<ValueType, Allocator>& OutValues, bool bMaintainOrder = false) const
    {
        Super::Read(InKey, [&InKey, &OutValues, bMaintainOrder](uint32 Hash, const typename Super::MapType& Map) { Map.MultiFind(InKey, OutValues, bMaintainOrder); });
    }
};