TypedElementOwnerStore.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "CoreMinimal.h"
#include "Misc/ScopeLock.h"
#include "Misc/TransactionallySafeCriticalSection.h"
#include "Elements/Framework/TypedElementHandle.h"
 
template <typename ElementOwnerType>
struct TTypedElementOwnerScopedAccessImpl
{
public:
    explicit TTypedElementOwnerScopedAccessImpl(FTransactionallySafeCriticalSection* InSynchObject)
        : SynchObject(InSynchObject)
    {
        check(SynchObject);
        SynchObject->Lock();
    }
 
    ~TTypedElementOwnerScopedAccessImpl()
    {
        if (SynchObject)
        {
            SynchObject->Unlock();
        }
    }
 
    TTypedElementOwnerScopedAccessImpl(const TTypedElementOwnerScopedAccessImpl&) = delete;
    TTypedElementOwnerScopedAccessImpl& operator=(const TTypedElementOwnerScopedAccessImpl&) = delete;
 
    TTypedElementOwnerScopedAccessImpl(TTypedElementOwnerScopedAccessImpl&& InOther)
        : SynchObject(InOther.SynchObject)
        , ElementOwner(InOther.ElementOwner)
    {
        InOther.SynchObject = nullptr;
        InOther.ElementOwner = nullptr;
    }
 
    TTypedElementOwnerScopedAccessImpl& operator=(TTypedElementOwnerScopedAccessImpl&&) = delete;
 
    inline explicit operator bool() const
    {
        return IsSet();
    }
 
    inline bool IsSet() const
    {
        return ElementOwner != nullptr;
    }
 
    inline ElementOwnerType& operator*() const
    {
        check(ElementOwner);
        return *ElementOwner;
    }
 
    inline ElementOwnerType* operator->() const
    {
        check(ElementOwner);
        return ElementOwner;
    }
 
    void Private_SetElementOwner(ElementOwnerType* InElementOwner)
    {
        checkf(!ElementOwner, TEXT("Element owner was already set!"));
        ElementOwner = InElementOwner;
 
        // If we didn't find an element then we can unlock immediately
        // Otherwise we hold the lock for as long as the element is being read to avoid a map write from invalidating the reference from reallocation
        check(SynchObject);
        if (!ElementOwner)
        {
            SynchObject->Unlock();
            SynchObject = nullptr;
        }
    }
 
private:
    FTransactionallySafeCriticalSection* SynchObject = nullptr;
    ElementOwnerType* ElementOwner = nullptr;
};
 
template <typename ElementDataType>
using TTypedElementOwnerScopedAccess = TTypedElementOwnerScopedAccessImpl<const TTypedElementOwner<ElementDataType>>;
 
template <typename ElementDataType>
using TTypedElementOwnerScopedMutableAccess = TTypedElementOwnerScopedAccessImpl<TTypedElementOwner<ElementDataType>>;
 
template <typename KeyDataType>
inline FString GetTypedElementOwnerStoreKeyDebugString(const KeyDataType& InKey)
{
    return TEXT("<GetTypedElementOwnerStoreKeyDebugString not implemented for key type>");
}
 
inline FString GetTypedElementOwnerStoreKeyDebugString(const UObject* InKey)
{
    return InKey->GetPathName();
}
 
template <typename ElementDataType, typename KeyDataType = const void*>
struct TTypedElementOwnerStore
{
public:
    static_assert(TNumericLimits<int32>::Max() >= TypedHandleMaxElementId, "TTypedElementOwnerStore internally uses TMap so cannot store TypedHandleMaxElementId! Consider making this container 64-bit aware, or explicitly remove this compile time check.");
 
    TTypedElementOwnerStore() = default;
 
    ~TTypedElementOwnerStore()
    {
        checkf(ElementOwnerMap.Num() == 0, TEXT("Element owners were still registered during destruction! This will leak elements, and you should unregister and destroy all elements prior to destruction!"));
    }
 
    TTypedElementOwnerScopedAccess<ElementDataType> RegisterElementOwner(const KeyDataType& InKey, TTypedElementOwner<ElementDataType>&& InElementOwner)
    {
        TTypedElementOwnerScopedAccess<ElementDataType> ScopedAccess(&ElementOwnerMapCS);
        checkf(!ElementOwnerMap.Contains(InKey), TEXT("Element owner has already been registered for this key (%s)! This will leak elements, and you should unregister and destroy the old element owner for this instance before adding a new one!"), *GetTypedElementOwnerStoreKeyDebugString(InKey));
        checkf(InElementOwner, TEXT("Element owner passed to RegisterElementOwner was null for key (%s)! Element owners must be valid!"), *GetTypedElementOwnerStoreKeyDebugString(InKey));
        ScopedAccess.Private_SetElementOwner(&ElementOwnerMap.Add(InKey, MoveTemp(InElementOwner)));
        return ScopedAccess;
    }
 
    TTypedElementOwner<ElementDataType> UnregisterElementOwner(const KeyDataType& InKey)
    {
        TTypedElementOwner<ElementDataType> ElementOwner;
        {
            UE::TScopeLock Lock(ElementOwnerMapCS);
            ElementOwnerMap.RemoveAndCopyValue(InKey, ElementOwner);
        }
        return ElementOwner;
    }
 
    void UnregisterElementOwners(TFunctionRef<bool(const TTypedElementOwner<ElementDataType>&)> InShouldUnregisterElement, TFunctionRef<void(TTypedElementOwner<ElementDataType>&&)> InOnUnregisterElement)
    {
        UE::TScopeLock Lock(ElementOwnerMapCS);
        for (auto It = ElementOwnerMap.CreateIterator(); It; ++It)
        {
            if (InShouldUnregisterElement(It->Value))
            {
                InOnUnregisterElement(MoveTemp(It->Value));
                It.RemoveCurrent();
            }
        }
    }
 
    TTypedElementOwnerScopedAccess<ElementDataType> FindOrRegisterElementOwner(const KeyDataType& InKey, TFunctionRef<TTypedElementOwner<ElementDataType>()> InCreateElement)
    {
        TTypedElementOwnerScopedAccess<ElementDataType> ScopedAccess(&ElementOwnerMapCS);
        if (const TTypedElementOwner<ElementDataType>* ExistingElement = ElementOwnerMap.Find(InKey))
        {
            ScopedAccess.Private_SetElementOwner(ExistingElement);
        }
        else if (TTypedElementOwner<ElementDataType> NewElement = InCreateElement())
        {
            ScopedAccess.Private_SetElementOwner(&ElementOwnerMap.Add(InKey, MoveTemp(NewElement)));
        }
        return ScopedAccess;
    }
 
    TTypedElementOwnerScopedAccess<ElementDataType> FindElementOwner(const KeyDataType& InKey) const
    {
        TTypedElementOwnerScopedAccess<ElementDataType> ScopedAccess(&ElementOwnerMapCS);
        ScopedAccess.Private_SetElementOwner(ElementOwnerMap.Find(InKey));
        return ScopedAccess;
    }
 
    TTypedElementOwnerScopedMutableAccess<ElementDataType> FindMutableElementOwner(const KeyDataType& InKey)
    {
        TTypedElementOwnerScopedMutableAccess<ElementDataType> ScopedAccess(&ElementOwnerMapCS);
        ScopedAccess.Private_SetElementOwner(ElementOwnerMap.Find(InKey));
        return ScopedAccess;
    }
 
    bool ContainsElementOwner(const KeyDataType& InKey) const
    {
        UE::TScopeLock Lock(ElementOwnerMapCS);
        return ElementOwnerMap.Contains(InKey);
    }
 
private:
    TMap<KeyDataType, TTypedElementOwner<ElementDataType>> ElementOwnerMap;
    mutable FTransactionallySafeCriticalSection ElementOwnerMapCS;
};