MovieScenePreAnimatedStateStorage.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "Containers/Map.h"
#include "Containers/SortedMap.h"
#include "Containers/SparseArray.h"
#include "CoreTypes.h"
#include "EntitySystem/BuiltInComponentTypes.h"
#include "EntitySystem/EntityAllocationIterator.h"
#include "EntitySystem/MovieSceneComponentPtr.h"
#include "EntitySystem/MovieSceneComponentTypeInfo.h"
#include "EntitySystem/MovieSceneEntityIDs.h"
#include "EntitySystem/MovieSceneEntityRange.h"
#include "EntitySystem/MovieSceneEntitySystemLinker.h"
#include "EntitySystem/MovieSceneEntitySystemTask.h"
#include "EntitySystem/MovieSceneEntitySystemTypes.h"
#include "Evaluation/PreAnimatedState/IMovieScenePreAnimatedStorage.h"
#include "Evaluation/PreAnimatedState/MovieScenePreAnimatedCaptureSources.h"
#include "Evaluation/PreAnimatedState/MovieScenePreAnimatedStateExtension.h"
#include "Evaluation/PreAnimatedState/MovieScenePreAnimatedStateTypes.h"
#include "Evaluation/PreAnimatedState/MovieScenePreAnimatedStorageID.h"
#include "Evaluation/PreAnimatedState/MovieSceneRestoreStateParams.h"
#include "Misc/AssertionMacros.h"
#include "Templates/Casts.h"
#include "Templates/SharedPointer.h"
#include "Templates/UnrealTemplate.h"
#include "Templates/UnrealTypeTraits.h"
#include "UObject/ObjectKey.h"
 
class FReferenceCollector;
class UObject;
 
namespace UE
{
namespace MovieScene
{
 
struct FPreAnimatedObjectGroupManager;
struct FRestoreStateParams;
struct FRootInstanceHandle;
 
enum class EPreAnimatedCaptureSourceTracking
{
    CacheIfTracked,
    AlwaysCache,
};
 
template<typename... InputTypes>
struct TPreAnimatedStateTaskParams
{
    FEntityComponentFilter AdditionalFilter;
    EPreAnimatedCaptureSourceTracking TrackingMode;
 
    TPreAnimatedStateTaskParams()
    {
        TrackingMode = EPreAnimatedCaptureSourceTracking::CacheIfTracked;
        AdditionalFilter.All({ FBuiltInComponentTypes::Get()->Tags.NeedsLink });
    }
};
 
struct FPreAnimatedStateTraits
{
    enum
    {
        NeedsInitialize = false,
        SupportsGrouping = false,
        SupportsReplaceObject = false,
    };
};
 
template<typename StorageTraits>
struct TPreAnimatedStateStorage : IPreAnimatedStorage
{
    using KeyType     = typename StorageTraits::KeyType;
    using StorageType = typename StorageTraits::StorageType;
 
    struct IRestoreMask
    {
        virtual ~IRestoreMask(){}
 
        virtual bool CanRestore(const KeyType& InKey) const = 0;
    };
 
    TPreAnimatedStateStorage()
    {}
 
    TPreAnimatedStateStorage(StorageTraits&& InTraits)
        : Traits(MoveTemp(InTraits))
    {}
 
    TPreAnimatedStateStorage(const TPreAnimatedStateStorage&) = delete;
    TPreAnimatedStateStorage& operator=(const TPreAnimatedStateStorage&) = delete;
 
public:
 
    void Initialize(FPreAnimatedStorageID InStorageID, FPreAnimatedStateExtension* InParentExtension) override
    {
        ParentExtension = InParentExtension;
        StorageID = InStorageID;
 
        if constexpr (StorageTraits::NeedsInitialize)
        {
            this->Traits.Initialize(InStorageID, InParentExtension);
        }
    }
 
    FPreAnimatedStorageID GetStorageType() const override
    {
        return StorageID;
    }
 
    EPreAnimatedStorageRequirement RestorePreAnimatedStateStorage(FPreAnimatedStorageIndex StorageIndex, EPreAnimatedStorageRequirement SourceRequirement, EPreAnimatedStorageRequirement TargetRequirement, const FRestoreStateParams& Params) override
    {
        if (RestoreMask)
        {
            if (!RestoreMask->CanRestore(PreAnimatedStorage[StorageIndex].Key))
            {
                return EPreAnimatedStorageRequirement::NoChange;
            }
        }
 
        if (SourceRequirement == EPreAnimatedStorageRequirement::Persistent)
        {
            // Restoring global state
            if (TargetRequirement == EPreAnimatedStorageRequirement::None)
            {
                FCachedData CachedData = MoveTemp(PreAnimatedStorage[StorageIndex]);
 
                KeyToStorageIndex.Remove(CachedData.Key);
                PreAnimatedStorage.RemoveAt(StorageIndex, 1);
                TransientPreAnimatedStorage.Remove(StorageIndex);
 
                if (CachedData.bInitialized)
                {
                    Traits.RestorePreAnimatedValue(CachedData.Key, CachedData.Value, Params);
                }
                return EPreAnimatedStorageRequirement::None;
            }
            else
            {
                ensure(TargetRequirement == EPreAnimatedStorageRequirement::NoChange);
 
                FCachedData& CachedData = PreAnimatedStorage[StorageIndex];
                if (CachedData.bInitialized)
                {
                    Traits.RestorePreAnimatedValue(CachedData.Key, CachedData.Value, Params);
                }
 
                return EPreAnimatedStorageRequirement::NoChange;
            }
        }
 
        ensure(SourceRequirement == EPreAnimatedStorageRequirement::Transient);
 
        // Always restore from the transient storage if available
        if (StorageType* CachedData = TransientPreAnimatedStorage.Find(StorageIndex))
        {
            Traits.RestorePreAnimatedValue(PreAnimatedStorage[StorageIndex].Key, *CachedData, Params);
 
            TransientPreAnimatedStorage.Remove(StorageIndex);
 
            return EPreAnimatedStorageRequirement::Persistent;
        }
 
        if (TargetRequirement == EPreAnimatedStorageRequirement::None)
        {
            FCachedData& ActualValue = PreAnimatedStorage[StorageIndex];
            if (ActualValue.bPersistent)
            {
                if (ActualValue.bInitialized)
                {
                    Traits.RestorePreAnimatedValue(ActualValue.Key, ActualValue.Value, Params);
                }
                return EPreAnimatedStorageRequirement::Persistent;
            }
 
            FCachedData Tmp = MoveTemp(PreAnimatedStorage[StorageIndex]);
 
            KeyToStorageIndex.Remove(Tmp.Key);
            PreAnimatedStorage.RemoveAt(StorageIndex, 1);
            TransientPreAnimatedStorage.Remove(StorageIndex);
 
            if (Tmp.bInitialized)
            {
                Traits.RestorePreAnimatedValue(Tmp.Key, Tmp.Value, Params);
            }
 
            return EPreAnimatedStorageRequirement::None;
        }
 
        if (TargetRequirement == EPreAnimatedStorageRequirement::Persistent)
        {
            // Restore the value but keep the value cached
            FCachedData& PersistentData = PreAnimatedStorage[StorageIndex];
 
            if (PersistentData.bInitialized)
            {
                PersistentData.bPersistent = true;
                Traits.RestorePreAnimatedValue(PersistentData.Key, PersistentData.Value, Params);
            }
        }
 
        return EPreAnimatedStorageRequirement::Persistent;
    }
 
    EPreAnimatedStorageRequirement DiscardPreAnimatedStateStorage(FPreAnimatedStorageIndex StorageIndex, EPreAnimatedStorageRequirement SourceRequirement) override
    {
        if (RestoreMask)
        {
            if (!RestoreMask->CanRestore(PreAnimatedStorage[StorageIndex].Key))
            {
                return EPreAnimatedStorageRequirement::NoChange;
            }
        }
 
        if (SourceRequirement == EPreAnimatedStorageRequirement::Persistent)
        {
            KeyType Key = PreAnimatedStorage[StorageIndex].Key;
 
            KeyToStorageIndex.Remove(Key);
            PreAnimatedStorage.RemoveAt(StorageIndex, 1);
            TransientPreAnimatedStorage.Remove(StorageIndex);
 
            return EPreAnimatedStorageRequirement::None;
        }
        else
        {
            ensure(SourceRequirement == EPreAnimatedStorageRequirement::Transient);
            const int32 NumTransients = TransientPreAnimatedStorage.Remove(StorageIndex);
            if (NumTransients == 0)
            {
                PreAnimatedStorage[StorageIndex].bPersistent = true;
            }
            return EPreAnimatedStorageRequirement::Persistent;
        }
    }
 
    void SetRestoreMask(const IRestoreMask* InRestoreMask)
    {
        RestoreMask = InRestoreMask;
    }
 
    void AddReferencedObjects(FReferenceCollector& ReferenceCollector) override
    {
        if constexpr (THasAddReferencedObjectForComponent<KeyType>::Value)
        {
            for (auto It = KeyToStorageIndex.CreateIterator(); It; ++It)
            {
                AddReferencedObjectForComponent(&ReferenceCollector, &It.Key());
            }
        }
        
        if constexpr (THasAddReferencedObjectForComponent<KeyType>::Value || THasAddReferencedObjectForComponent<StorageType>::Value)
        {
            for (FCachedData& CachedData : PreAnimatedStorage)
            {
                AddReferencedObjectForComponent(&ReferenceCollector, &CachedData.Key);
                AddReferencedObjectForComponent(&ReferenceCollector, &CachedData.Value);
            }
        }
        if constexpr (THasAddReferencedObjectForComponent<StorageType>::Value)
        {
            for (auto It = TransientPreAnimatedStorage.CreateIterator(); It; ++It)
            {
                AddReferencedObjectForComponent(&ReferenceCollector, &It.Value());
            }
        }
    }
 
    FPreAnimatedStorageIndex GetOrCreateStorageIndex(const KeyType& InKey)
    {
        FPreAnimatedStorageIndex Index = KeyToStorageIndex.FindRef(InKey);
        if (!Index)
        {
            Index = PreAnimatedStorage.Add(FCachedData{InKey});
            KeyToStorageIndex.Add(InKey, Index);
        }
 
        return Index;
    }
 
    FPreAnimatedStorageIndex FindStorageIndex(const KeyType& InKey) const
    {
        return KeyToStorageIndex.FindRef(InKey);
    }
 
    void AssignPreAnimatedValue(FPreAnimatedStorageIndex StorageIndex, EPreAnimatedStorageRequirement StorageRequirement, StorageType&& InNewValue)
    {
        check(StorageIndex);
 
        FCachedData& CachedData = PreAnimatedStorage[StorageIndex.Value];
 
        if (StorageRequirement == EPreAnimatedStorageRequirement::Persistent)
        {
            ensure(!CachedData.bInitialized);
            CachedData.Value = MoveTemp(InNewValue);
            CachedData.bPersistent = true;
            CachedData.bInitialized = true;
        }
        else if (StorageRequirement == EPreAnimatedStorageRequirement::Transient)
        {
            ensure(!CachedData.bInitialized || !TransientPreAnimatedStorage.Contains(StorageIndex));
 
            // Assign the transient value
            if (!CachedData.bInitialized)
            {
                CachedData.Value = MoveTemp(InNewValue);
                CachedData.bInitialized = true;
            }
            else
            {
                TransientPreAnimatedStorage.Add(StorageIndex, MoveTemp(InNewValue));
            }
        }
    }
 
    bool IsStorageRequirementSatisfied(FPreAnimatedStorageIndex StorageIndex, EPreAnimatedStorageRequirement StorageRequirement) const
    {
        check(StorageIndex);
 
        const FCachedData& CachedData = PreAnimatedStorage[StorageIndex.Value];
 
        if (StorageRequirement == EPreAnimatedStorageRequirement::Persistent)
        {
            return CachedData.bInitialized;
        }
        else if (StorageRequirement == EPreAnimatedStorageRequirement::Transient)
        {
            return (CachedData.bInitialized && CachedData.bPersistent == false)
                || TransientPreAnimatedStorage.Contains(StorageIndex);
        }
 
        return true;
    }
 
    template<typename... KeyArgs>
    FPreAnimatedStateEntry FindEntry(KeyArgs&&... InKeyArgs)
    {
        KeyType Key{ Forward<KeyArgs>(InKeyArgs)... };
 
        FPreAnimatedStorageGroupHandle GroupHandle;
        if constexpr (StorageTraits::SupportsGrouping)
        {
            GroupHandle = this->Traits.FindGroup(this->Traits.ResolveComponent(InKeyArgs)...);
        }
 
        FPreAnimatedStorageIndex StorageIndex = this->FindStorageIndex(Key);
 
        return FPreAnimatedStateEntry{ GroupHandle, FPreAnimatedStateCachedValueHandle{ this->StorageID, StorageIndex } };
    }
 
    template<typename... KeyArgs>
    FPreAnimatedStateEntry MakeEntry(KeyArgs&&... InKeyArgs)
    {
        KeyType Key{ Forward<KeyArgs>(InKeyArgs)... };
 
        FPreAnimatedStorageGroupHandle GroupHandle;
        if constexpr(StorageTraits::SupportsGrouping)
        {
            GroupHandle = this->Traits.MakeGroup(this->Traits.ResolveComponent(InKeyArgs)...);
        }
 
        FPreAnimatedStorageIndex StorageIndex = this->GetOrCreateStorageIndex(Key);
 
        return FPreAnimatedStateEntry{ GroupHandle, FPreAnimatedStateCachedValueHandle{ this->StorageID, StorageIndex } };
    }
 
    bool IsCapturingGlobalState() const
    {
        return this->ParentExtension->IsCapturingGlobalState();
    }
 
    void ForciblyPersistStorage(FPreAnimatedStorageIndex StorageIndex)
    {
        check(StorageIndex);
        PreAnimatedStorage[StorageIndex.Value].bPersistent = true;
    }
 
    bool IsStorageInitialized(FPreAnimatedStorageIndex StorageIndex) const
    {
        return StorageIndex && (PreAnimatedStorage[StorageIndex.Value].bInitialized || TransientPreAnimatedStorage.Contains(StorageIndex));
    }
 
    bool HasEverAnimated(FPreAnimatedStorageIndex StorageIndex) const
    {
        return StorageIndex && PreAnimatedStorage[StorageIndex.Value].bInitialized;
    }
 
    const KeyType& GetKey(FPreAnimatedStorageIndex StorageIndex) const
    {
        return PreAnimatedStorage[StorageIndex].Key;
    }
 
    void ReplaceKey(FPreAnimatedStorageIndex StorageIndex, const KeyType& NewKey)
    {
        KeyType OldKey = PreAnimatedStorage[StorageIndex].Key;
        PreAnimatedStorage[StorageIndex].Key = NewKey;
 
        KeyToStorageIndex.Remove(OldKey);
        KeyToStorageIndex.Add(NewKey, StorageIndex);
    }
 
    const StorageType& GetCachedValue(FPreAnimatedStorageIndex StorageIndex) const
    {
        static const StorageType DefaultValue = StorageType();
 
        if (ensure(PreAnimatedStorage.IsValidIndex(StorageIndex.Value)))
        {
            const FCachedData& CachedData = PreAnimatedStorage[StorageIndex.Value];
            if (CachedData.bInitialized)
            {
                return CachedData.Value;
            }
        }
        return DefaultValue;
    }
 
    template<typename... ContributorTypes>
    void BeginTrackingEntities(UMovieSceneEntitySystemLinker* Linker, TComponentTypeID<ContributorTypes>... InComponentTypes)
    {
        BeginTrackingEntitiesTask(Linker, TPreAnimatedStateTaskParams<ContributorTypes...>(), InComponentTypes...);
    }
 
    template<typename TaskType, typename... ContributorTypes>
    void BeginTrackingEntitiesTask(UMovieSceneEntitySystemLinker* Linker, const TaskType& InParams, TComponentTypeID<ContributorTypes>... InComponentTypes)
    {
        auto VisitAllocation = [this, InParams](FEntityAllocationIteratorItem Item, TRead<FMovieSceneEntityID> EntityIDs, TRead<FRootInstanceHandle> RootInstanceHandles, TRead<ContributorTypes>... Inputs)
        {
            FPreAnimatedTrackerParams Params(Item);
 
            const int32 Num = Params.Num;
            const bool  bWantsRestore = Params.bWantsRestoreState;
 
            if (!this->ParentExtension->IsCapturingGlobalState() && !bWantsRestore)
            {
                return;
            }
 
            FPreAnimatedEntityCaptureSource* EntityMetaData = this->ParentExtension->GetOrCreateEntityMetaData();
            for (int32 Index = 0; Index < Num; ++Index)
            {
                FPreAnimatedStateEntry Entry = MakeEntry(Inputs[Index]...);
                EntityMetaData->BeginTrackingEntity(Entry, EntityIDs[Index], RootInstanceHandles[Index], bWantsRestore);
            }
        };
 
        FBuiltInComponentTypes* BuiltInComponents = FBuiltInComponentTypes::Get();
 
        FComponentMask ExcludeMask({ BuiltInComponents->Tags.NeedsUnlink, BuiltInComponents->Tags.Finished, BuiltInComponents->Tags.Ignored });
 
        FEntityTaskBuilder()
        .ReadEntityIDs()
        .Read(BuiltInComponents->RootInstanceHandle)
        .ReadAllOf(InComponentTypes...)
        .FilterNone(ExcludeMask)
        .CombineFilter(InParams.AdditionalFilter)
        .Iterate_PerAllocation(&Linker->EntityManager, VisitAllocation);
    }
 
    template<typename... ContributorTypes>
    void BeginTrackingEntity(FMovieSceneEntityID EntityID, const bool bWantsRestoreState, FRootInstanceHandle RootInstanceHandle, ContributorTypes... InComponents)
    {
        if (!this->ParentExtension->IsCapturingGlobalState() && !bWantsRestoreState)
        {
            return;
        }
 
        TPreAnimatedStateTaskParams<ContributorTypes...> Params;
 
        FPreAnimatedStateEntry Entry = MakeEntry(InComponents...);
        FPreAnimatedEntityCaptureSource* EntityMetaData = this->ParentExtension->GetOrCreateEntityMetaData();
        EntityMetaData->BeginTrackingEntity(Entry, EntityID, RootInstanceHandle, bWantsRestoreState);
    }
 
    template<typename... ContributorTypes>
    void CachePreAnimatedValues(UMovieSceneEntitySystemLinker* Linker, TComponentTypeID<ContributorTypes>... InComponentTypes)
    {
        CachePreAnimatedValuesTask(Linker, TPreAnimatedStateTaskParams<ContributorTypes...>(), InComponentTypes...);
    }
 
    template<typename TaskType, typename... ContributorTypes>
    void CachePreAnimatedValuesTask(UMovieSceneEntitySystemLinker* Linker, const TaskType& InParams, TComponentTypeID<ContributorTypes>... InComponentTypes)
    {
        auto VisitAllocation = [this, &InParams](FEntityAllocationIteratorItem Item, TRead<ContributorTypes>... Values)
        {
            const int32 Num = Item.GetAllocation()->Num();
            for (int32 Index = 0; Index < Num; ++Index)
            {
                if (!ShouldTrackCaptureSource(InParams.TrackingMode, Values[Index]...))
                {
                    continue;
                }
 
                FPreAnimatedStateEntry Entry = MakeEntry(Values[Index]...);
                TrackCaptureSource(Entry, InParams.TrackingMode);
                
                EPreAnimatedStorageRequirement StorageRequirement = this->ParentExtension->GetStorageRequirement(Entry);
                if (!this->IsStorageRequirementSatisfied(Entry.ValueHandle.StorageIndex, StorageRequirement))
                {
                    StorageType NewValue = this->Traits.CachePreAnimatedValue(this->Traits.ResolveComponent(Values[Index])...);
                    this->AssignPreAnimatedValue(Entry.ValueHandle.StorageIndex, StorageRequirement, MoveTemp(NewValue));
                }
            }
        };
 
        FEntityTaskBuilder()
        .ReadAllOf(InComponentTypes...)
        .CombineFilter(InParams.AdditionalFilter)
        .Iterate_PerAllocation(&Linker->EntityManager, VisitAllocation);
    }
 
    template<typename... ContributorTypes>
    void CachePreAnimatedValue(ContributorTypes... Values)
    {
        CacheTrackedPreAnimatedValue(EPreAnimatedCaptureSourceTracking::CacheIfTracked, Values...);
    }
 
    template<typename... ContributorTypes>
    void CacheTrackedPreAnimatedValue(EPreAnimatedCaptureSourceTracking TrackingMode, ContributorTypes... Values)
    {
        if (ShouldTrackCaptureSource(TrackingMode, Values...))
        {
            FPreAnimatedStateEntry Entry = MakeEntry(Values...);
 
            TrackCaptureSource(Entry, TrackingMode);
 
            EPreAnimatedStorageRequirement StorageRequirement = this->ParentExtension->GetStorageRequirement(Entry);
            if (!IsStorageRequirementSatisfied(Entry.ValueHandle.StorageIndex, StorageRequirement))
            {
                StorageType NewValue = this->Traits.CachePreAnimatedValue(this->Traits.ResolveComponent(Values)...);
                this->AssignPreAnimatedValue(Entry.ValueHandle.StorageIndex, StorageRequirement, MoveTemp(NewValue));
            }
        }
    }
 
    template<typename OnCacheValue /* StorageType(const KeyType&) */>
    void CachePreAnimatedValue(const KeyType& InKey, OnCacheValue&& CacheCallback, EPreAnimatedCaptureSourceTracking TrackingMode = EPreAnimatedCaptureSourceTracking::CacheIfTracked)
    {
        static_assert(StorageTraits::SupportsGrouping == false, "Grouped pre-animated state requires passing a group handle");
        CachePreAnimatedValue(FPreAnimatedStorageGroupHandle(), InKey, Forward<OnCacheValue>(CacheCallback), TrackingMode);
    }
 
    template<typename OnCacheValue /* StorageType(const KeyType&) */>
    void CachePreAnimatedValue(FPreAnimatedStorageGroupHandle GroupHandle, const KeyType& InKey, OnCacheValue&& CacheCallback, EPreAnimatedCaptureSourceTracking TrackingMode = EPreAnimatedCaptureSourceTracking::CacheIfTracked)
    {
        ensureMsgf(GroupHandle.IsValid() || !StorageTraits::SupportsGrouping, TEXT("The group handle must be valid for pre-animated state that supports grouping, and invalid if not"));
 
        if (!ShouldTrackCaptureSource(TrackingMode, InKey))
        {
            return;
        }
        // Find the storage index for the specific key we're animating
        FPreAnimatedStorageIndex StorageIndex = GetOrCreateStorageIndex(InKey);
        FPreAnimatedStateEntry Entry{ GroupHandle, FPreAnimatedStateCachedValueHandle{ this->StorageID, StorageIndex } };
        TrackCaptureSource(Entry, TrackingMode);
 
        EPreAnimatedStorageRequirement StorageRequirement = this->ParentExtension->GetStorageRequirement(Entry);
        if (!IsStorageRequirementSatisfied(Entry.ValueHandle.StorageIndex, StorageRequirement))
        {
            StorageType NewValue = CacheCallback(InKey);
            AssignPreAnimatedValue(StorageIndex, StorageRequirement, MoveTemp(NewValue));
        }
    }
 
    template<typename... ContributorTypes>
    void BeginTrackingAndCachePreAnimatedValues(UMovieSceneEntitySystemLinker* Linker, TComponentTypeID<ContributorTypes>... InComponentTypes)
    {
        BeginTrackingAndCachePreAnimatedValuesTask(Linker, TPreAnimatedStateTaskParams<ContributorTypes...>(), InComponentTypes...);
    }
 
    template<typename TaskType, typename... ContributorTypes>
    void BeginTrackingAndCachePreAnimatedValuesTask(UMovieSceneEntitySystemLinker* Linker, const TaskType& InParams, TComponentTypeID<ContributorTypes>... InComponentTypes)
    {
        auto VisitAllocation = [this, InParams](FEntityAllocationIteratorItem Item, TRead<FMovieSceneEntityID> EntityIDs, TRead<FRootInstanceHandle> RootInstanceHandles, TRead<ContributorTypes>... Inputs)
        {
            FPreAnimatedTrackerParams Params(Item);
 
            const int32 Num = Params.Num;
            const bool  bWantsRestore = Params.bWantsRestoreState;
 
            if (!this->ParentExtension->IsCapturingGlobalState() && !bWantsRestore)
            {
                return;
            }
 
            FPreAnimatedEntityCaptureSource* EntityMetaData = this->ParentExtension->GetOrCreateEntityMetaData();
            for (int32 Index = 0; Index < Num; ++Index)
            {
                FPreAnimatedStateEntry Entry = MakeEntry(Inputs[Index]...);
 
                EntityMetaData->BeginTrackingEntity(Entry, EntityIDs[Index], RootInstanceHandles[Index], bWantsRestore);
 
                EPreAnimatedStorageRequirement StorageRequirement = this->ParentExtension->GetStorageRequirement(Entry);
                if (!this->IsStorageRequirementSatisfied(Entry.ValueHandle.StorageIndex, StorageRequirement))
                {
                    StorageType NewValue = this->Traits.CachePreAnimatedValue(this->Traits.ResolveComponent(Inputs[Index])...);
                    this->AssignPreAnimatedValue(Entry.ValueHandle.StorageIndex, StorageRequirement, MoveTemp(NewValue));
                }
            }
        };
 
        FBuiltInComponentTypes* BuiltInComponents = FBuiltInComponentTypes::Get();
 
        FEntityTaskBuilder()
        .ReadEntityIDs()
        .Read(BuiltInComponents->RootInstanceHandle)
        .ReadAllOf(InComponentTypes...)
        .CombineFilter(InParams.AdditionalFilter)
        .Iterate_PerAllocation(&Linker->EntityManager, VisitAllocation);
    }
 
    void OnObjectReplaced(FPreAnimatedStorageIndex StorageIndex, const FObjectKey& OldObject, const FObjectKey& NewObject) override
    {
        KeyType ExistingKey = this->GetKey(StorageIndex);
        if constexpr(StorageTraits::SupportsReplaceObject)
        {
            this->Traits.ReplaceObject(ExistingKey, NewObject);
        }
        this->ReplaceKey(StorageIndex, ExistingKey);
    }
 
protected:
 
    template<typename... KeyArgs>
    bool ShouldTrackCaptureSource(EPreAnimatedCaptureSourceTracking TrackingMode, KeyArgs&&... InKeyArgs)
    {
        switch (TrackingMode)
        {
        case EPreAnimatedCaptureSourceTracking::CacheIfTracked:
        {
            // If we're capturing global state we always track changes
            if (this->ParentExtension->IsCapturingGlobalState())
            {
                return true;
            }
 
            // Only cache the value if an entry already exists, and if metadata already exists.
            // (ie, something is actively animating this and wants restore state)
            if (FPreAnimatedStateEntry Entry = FindEntry(InKeyArgs...))
            {
                return ensureMsgf(this->ParentExtension->MetaDataExists(Entry), TEXT("PreAnimatedStateEntry has allocated storage but no metadata exists to track it."));
            }
            return false;
        }
        break;
 
        case EPreAnimatedCaptureSourceTracking::AlwaysCache:
        {
            // Always cache the source meta data. If there is no capture source this calls EnsureMetaData as above
            return true;
        }
        break;
 
        default:
            ensureMsgf(false, TEXT("Unsupported tracking mode, no pre-animated state caching will occur!"));
            return false;
        }
    }
 
    void TrackCaptureSource(const FPreAnimatedStateEntry& Entry, EPreAnimatedCaptureSourceTracking TrackingMode)
    {
        switch (TrackingMode)
        {
        case EPreAnimatedCaptureSourceTracking::CacheIfTracked:
            {
                // If we're capturing global state we always track changes
                if (this->ParentExtension->IsCapturingGlobalState())
                {
                    this->ParentExtension->EnsureMetaData(Entry);
                    return;
                }
 
                // Only cache the value if tracking meta-data exists for this entry
                // (ie, something is actively animating this and wants restore state)
                ensureMsgf(Entry.IsValid() && this->ParentExtension->MetaDataExists(Entry), TEXT("PreAnimatedStateEntry has allocated storage but no metadata exists to track it."));
            }
            break;
 
        case EPreAnimatedCaptureSourceTracking::AlwaysCache:
            {
                // Always cache the source meta data. If there is no capture source this calls EnsureMetaData as above
                this->ParentExtension->AddSourceMetaData(Entry);
                return;
            }
            break;
 
        default:
            ensureMsgf(false, TEXT("Unsupported tracking mode, no pre-animated state caching will occur!"));
        }
    }
 
    struct FCachedData
    {
        FCachedData()
            : bInitialized(false)
            , bPersistent(false)
        {}
 
        FCachedData(const KeyType& InKey)
            : Key(InKey)
            , bInitialized(false)
            , bPersistent(false)
        {}
 
        FCachedData(const KeyType& InKey, StorageType&& InValue)
            : Key(InKey)
            , Value(MoveTemp(InValue))
            , bInitialized(true)
            , bPersistent(false)
        {}
 
        KeyType Key;
        StorageType Value;
        bool bInitialized : 1;
        bool bPersistent : 1;
    };
 
    TMap<KeyType, FPreAnimatedStorageIndex> KeyToStorageIndex;
 
    TSparseArray<FCachedData> PreAnimatedStorage;
 
    TSortedMap<FPreAnimatedStorageIndex, StorageType> TransientPreAnimatedStorage;
 
    FPreAnimatedStateExtension* ParentExtension = nullptr;
 
    const IRestoreMask* RestoreMask = nullptr;
 
    FPreAnimatedStorageID StorageID;
 
public:
 
    StorageTraits Traits;
};
 
 
template<typename InKeyType, typename InStorageType>
struct TSimplePreAnimatedStateTraits : FPreAnimatedStateTraits
{
    using KeyType = InKeyType;
    using StorageType = InStorageType;
 
    void RestorePreAnimatedValue(const KeyType& Key, StorageType& SavedValue, const FRestoreStateParams& Params)
    {
        SavedValue.RestoreState(Key, Params);
    }
};
 
 
template<typename KeyType, typename StorageType>
struct TSimplePreAnimatedStateStorage : TPreAnimatedStateStorage<TSimplePreAnimatedStateTraits<KeyType, StorageType>>
{
};
 
} // namespace MovieScene
} // namespace UE