MovieSceneEntityGroupingSystem.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "Containers/Map.h"
#include "Containers/SparseArray.h"
#include "CoreTypes.h"
#include "EntitySystem/BuiltInComponentTypes.h"
#include "EntitySystem/MovieSceneComponentTypeInfo.h"
#include "EntitySystem/MovieSceneEntityIDs.h"
#include "EntitySystem/MovieSceneEntitySystem.h"
#include "EntitySystem/MovieSceneEntitySystemTask.h"
#include "EntitySystem/MovieSceneEntitySystemTypes.h"
#include "Templates/PointerIsConvertibleFromTo.h"
#include "UObject/ObjectMacros.h"
#include "UObject/UObjectGlobals.h"
 
#include "MovieSceneEntityGroupingSystem.generated.h"
 
#define UE_API MOVIESCENE_API
 
class UMovieSceneEntityGroupingSystem;
 
namespace UE::MovieScene
{
 
struct FAddGroupMutation;
struct FUpdateGroupsTask;
 
struct CBatchGroupingPolicy
{
    template <typename T>
    auto Requires(T& In) -> decltype(&T::InitializeGroupKeys);
};
 
struct FEntityGroupBuilder
{
    UE_API FEntityGroupBuilder(UMovieSceneEntityGroupingSystem* InOwner, FEntityGroupingPolicyKey InPolicyKey);
 
    FEntityGroupID MakeGroupID(int32 GroupIndex) const { return FEntityGroupID(PolicyKey, GroupIndex); }
    FEntityGroupID MakeInvalidGroupID() const { return FEntityGroupID(PolicyKey, INDEX_NONE); }
    UE_API void AddEntityToGroup(const FMovieSceneEntityID& InEntity, const FEntityGroupID& InNewGroupID);
    UE_API void RemoveEntityFromGroup(const FMovieSceneEntityID& InEntity, const FEntityGroupID& InPreviousGroupID);
 
    UE_API int32 AllocateGroupIndex();
 
    UE_API void ReportUsedGroupIndex(int32 GroupIndex);
 
private:
    UMovieSceneEntityGroupingSystem* Owner;
    FEntityGroupingPolicyKey PolicyKey;
};
 
struct IEntityGroupingHandler
{
    virtual ~IEntityGroupingHandler() {}
    virtual void ProcessAllocation(FEntityAllocationIteratorItem Item, FReadEntityIDs EntityIDs, TWrite<FEntityGroupID> GroupIDs, FEntityGroupBuilder* Builder) = 0;
    virtual void OnGroupIndexFreed(int32 InGroupIndex) {}
 
#if WITH_EDITOR
    virtual void OnObjectsReplaced(const TMap<UObject*, UObject*>& ReplacementMap) = 0;
#endif
};
 
template<typename GroupingPolicy, typename ComponentIndices, typename ...ComponentTypes>
struct TEntityGroupingHandlerImpl;
 
template<typename GroupingPolicy, typename ...ComponentTypes>
struct TEntityGroupingHandler 
    : TEntityGroupingHandlerImpl<GroupingPolicy, TMakeIntegerSequence<int, sizeof...(ComponentTypes)>, ComponentTypes...>
{
    TEntityGroupingHandler(GroupingPolicy&& InPolicy, TComponentTypeID<ComponentTypes>... InComponents)
        : TEntityGroupingHandlerImpl<GroupingPolicy, TMakeIntegerSequence<int, sizeof...(ComponentTypes)>, ComponentTypes...>(
                MoveTemp(InPolicy), InComponents...)
    {
    }
};
 
 
 
template<typename GroupingPolicy>
struct TEntityGroupingHandlerBase
{
    using GroupKeyType = typename GroupingPolicy::GroupKeyType;
 
    int32 GetOrAllocateGroupIndex(typename TCallTraits<GroupKeyType>::ParamType InGroupKey, FEntityGroupBuilder* Builder)
    {
        int32& GroupIndex = GroupKeyToIndex.FindOrAdd(InGroupKey, INDEX_NONE);
        if (GroupIndex == INDEX_NONE)
        {
            // This group key isn't known to us... let's allocate a new group index for it.
            // Try to find an available index first. Otherwise use a new high index.
            GroupIndex = Builder->AllocateGroupIndex();
            IndexToGroupKey.FindOrAdd(GroupIndex, InGroupKey);
        }
        else
        {
            // We know this group key, so we'll return the group index we already have
            // associated with it. We just need to "revive" it in case it was scheduled
            // for being freed.
            Builder->ReportUsedGroupIndex(GroupIndex);
        }
        return GroupIndex;
    }
 
protected:
 
    TMap<GroupKeyType, int32> GroupKeyToIndex;
    TMap<int32, GroupKeyType> IndexToGroupKey;
};
 
 
 
 
template<typename GroupingPolicy, int ...ComponentIndices, typename ...ComponentTypes>
struct TEntityGroupingHandlerImpl<GroupingPolicy, TIntegerSequence<int, ComponentIndices...>, ComponentTypes...> : TEntityGroupingHandlerBase<GroupingPolicy>, IEntityGroupingHandler
{
    using GroupKeyType = typename GroupingPolicy::GroupKeyType;
 
    GroupingPolicy Policy;
 
    TTuple<TComponentTypeID<ComponentTypes>...> Components;
 
    TEntityGroupingHandlerImpl(GroupingPolicy&& InPolicy, TComponentTypeID<ComponentTypes>... InComponents)
        : Policy(MoveTemp(InPolicy))
        , Components(InComponents...)
    {
    }
 
    virtual void ProcessAllocation(FEntityAllocationIteratorItem Item,  FReadEntityIDs EntityIDs, TWrite<FEntityGroupID> GroupIDs, FEntityGroupBuilder* Builder) override
    {
        const FBuiltInComponentTypes* BuiltInComponents = FBuiltInComponentTypes::Get();
        const FComponentMask& AllocationType = Item.GetAllocationType();
        const bool bNeedsLink = AllocationType.Contains(BuiltInComponents->Tags.NeedsLink);
        const bool bNeedsUnlink = AllocationType.Contains(BuiltInComponents->Tags.NeedsUnlink);
        ensure(bNeedsLink || bNeedsUnlink);
 
        if (bNeedsLink)
        {
            VisitLinkedEntities(Item, EntityIDs, GroupIDs, Builder);
        }
        else if (bNeedsUnlink)
        {
            VisitUnlinkedEntities(Item, EntityIDs, GroupIDs, Builder);
        }
    }
    
    void VisitLinkedEntities(FEntityAllocationIteratorItem Item, FReadEntityIDs EntityIDs, TWrite<FEntityGroupID> GroupIDs, FEntityGroupBuilder* Builder)
    {
        if constexpr (TModels_V<CBatchGroupingPolicy, GroupingPolicy>)
        {
            const FEntityAllocation* Allocation = Item.GetAllocation();
            Policy.InitializeGroupKeys(*this, Builder, Item, EntityIDs, GroupIDs, Allocation->ReadComponents(Components.template Get<ComponentIndices>())...);
        }
        else
        {
            const FEntityAllocation* Allocation = Item.GetAllocation();
            const int32 Num = Allocation->Num();
 
            const FEntityGroupID InvalidGroupID = Builder->MakeInvalidGroupID();
 
            TTuple<TComponentReader<ComponentTypes>...> ComponentReaders(
                    Allocation->ReadComponents(Components.template Get<ComponentIndices>())...);
 
            for (int32 Index = 0; Index < Num; ++Index)
            {
                GroupKeyType GroupKey;
                const bool bValidGroupKey = Policy.GetGroupKey(ComponentReaders.template Get<ComponentIndices>()[Index]..., GroupKey);
 
                const FMovieSceneEntityID EntityID(EntityIDs[Index]);
                FEntityGroupID& GroupID(GroupIDs[Index]);
 
                if (bValidGroupKey)
                {
                    // Find or create the appropriate group and put the entity in it.
                    int32 NewGroupIndex = this->GetOrAllocateGroupIndex(GroupKey, Builder);
                    FEntityGroupID NewGroupID = Builder->MakeGroupID(NewGroupIndex);
                    Builder->AddEntityToGroup(EntityID, NewGroupID);
                    GroupID = NewGroupID;
                }
                else
                {
                    // This entity doesn't belong to any group.
                    // Let's assign an invalid group ID that nonetheless has a valid policy key
                    // pointing to this grouping.
                    GroupID = InvalidGroupID;
                }
            }
        }
    }
 
    void VisitUnlinkedEntities(FEntityAllocationIteratorItem Item, FReadEntityIDs EntityIDs, TWrite<FEntityGroupID> GroupIDs, FEntityGroupBuilder* Builder)
    {
        const FEntityAllocation* Allocation = Item.GetAllocation();
        const int32 Num = Allocation->Num();
 
        for (int32 Index = 0; Index < Num; ++Index)
        {
            const FMovieSceneEntityID EntityID(EntityIDs[Index]);
            FEntityGroupID& GroupID(GroupIDs[Index]);
 
            if (GroupID.HasGroup())
            {
                Builder->RemoveEntityFromGroup(EntityID, GroupID);
                // Leave the GroupID on the entity so that downstream systems can use it to track
                // that this entity is leaving its group, but flag it so we don't re-free it.
                ensure(!EnumHasAllFlags(GroupID.Flags , EEntityGroupFlags::RemovedFromGroup));
                GroupID.Flags |= EEntityGroupFlags::RemovedFromGroup;
            }
        }
    }
 
    virtual void OnGroupIndexFreed(int32 GroupIndex) override
    {
        GroupKeyType Key = this->IndexToGroupKey.FindChecked(GroupIndex);
        ensure(this->IndexToGroupKey.Remove(GroupIndex) == 1);
        ensure(this->GroupKeyToIndex.Remove(Key) == 1);
    }
 
#if WITH_EDITOR
    virtual void OnObjectsReplaced(const TMap<UObject*, UObject*>& ReplacementMap) override
    {
        // Get a list of keys that contain replaced objects.
        TMap<GroupKeyType, GroupKeyType> ReplacedKeys;
        for (const TPair<GroupKeyType, int32>& Pair : this->GroupKeyToIndex)
        {
            GroupKeyType NewKey = Pair.Key;
            if (Policy.OnObjectsReplaced(NewKey, ReplacementMap))
            {
                ReplacedKeys.Add(Pair.Key, NewKey);
            }
        }
        // Replace the keys but keep the group indices.
        for (const TPair<GroupKeyType, GroupKeyType>& Pair : ReplacedKeys)
        {
            int32 GroupIndex;
            const bool bRemoved = this->GroupKeyToIndex.RemoveAndCopyValue(Pair.Key, GroupIndex);
            if (ensure(bRemoved))
            {
                this->GroupKeyToIndex.Add(Pair.Value, GroupIndex);
                this->IndexToGroupKey.Add(GroupIndex, Pair.Value);
            }
        }
    }
#endif
};
 
namespace Private
{
    template<typename T>
    bool ReplaceGroupKeyObjectElement(T&& InElem, const TMap<UObject*, UObject*>& ReplacementMap)
    {
        return false;
    }
 
    template<typename T>
    typename TEnableIf<TPointerIsConvertibleFromTo<T, UObject>::Value, bool>::Type ReplaceGroupKeyObjectElement(T& InOutElem, const TMap<UObject*, UObject*>& ReplacementMap)
    {
        UObject* CurrentObject = InOutElem;
        if (UObject* const* NewObject = ReplacementMap.Find(CurrentObject))
        {
            InOutElem = NewObject;
            return true;
        }
        return false;
    }
}
 
template<typename... ComponentTypes>
struct TTupleGroupingPolicy
{
    using GroupKeyType = TTuple<ComponentTypes...>;
 
    bool GetGroupKey(ComponentTypes... InComponents, GroupKeyType& OutGroupKey)
    {
        OutGroupKey = MakeTuple(InComponents...);
        return true;
    }
 
    bool OnObjectsReplaced(GroupKeyType& InOutKey, const TMap<UObject*, UObject*>& ReplacementMap)
    {
        bool bReplaced = false;
        VisitTupleElements([&ReplacementMap, &bReplaced](auto& InElem)
            {
                bReplaced |= Private::ReplaceGroupKeyObjectElement(InElem, ReplacementMap);
            },
            InOutKey);
        return bReplaced;
    }
};
 
} // namespace UE::MovieScene
 
 
UCLASS(MinimalAPI)
class UMovieSceneEntityGroupingSystem : public UMovieSceneEntitySystem
{
public:
 
    using FEntityGroupID = UE::MovieScene::FEntityGroupID;
    using FEntityGroupingPolicyKey = UE::MovieScene::FEntityGroupingPolicyKey;
    using IEntityGroupingHandler = UE::MovieScene::IEntityGroupingHandler;
    using FMovieSceneEntityID = UE::MovieScene::FMovieSceneEntityID;
 
    GENERATED_BODY()
 
    UE_API UMovieSceneEntityGroupingSystem(const FObjectInitializer& ObjInit);
 
    template<typename GroupingPolicy, typename ...ComponentTypes>
    FEntityGroupingPolicyKey AddGrouping(GroupingPolicy&& InPolicy, TComponentTypeID<ComponentTypes>... InComponents)
    {
        return AddGrouping(Forward<GroupingPolicy>(InPolicy), UE::MovieScene::FEntityComponentFilter(), InComponents...);
    }
 
    template<typename GroupingPolicy, typename ...ComponentTypes>
    FEntityGroupingPolicyKey AddGrouping(GroupingPolicy&& InPolicy, UE::MovieScene::FEntityComponentFilter&& InComponentFilter, TComponentTypeID<ComponentTypes>... InComponents)
    {
        using namespace UE::MovieScene;
 
        using NewGroupHandlerType = TEntityGroupingHandler<GroupingPolicy, ComponentTypes...>;
 
        static_assert(sizeof(NewGroupHandlerType) <= 256, "Handler type too big! Please increase the TInlineValue size.");
 
        const int32 HandlerIndex = GroupHandlers.Emplace();
        FEntityGroupingHandlerInfo& HandlerInfo = GroupHandlers[HandlerIndex];
        HandlerInfo.Handler.Emplace<NewGroupHandlerType>(MoveTemp(InPolicy), InComponents...);
        HandlerInfo.ComponentFilter = MoveTemp(InComponentFilter);
        HandlerInfo.ComponentFilter.All({ InComponents... });
 
        FEntityGroupingPolicyKey NewPolicyKey{ HandlerIndex };
        return NewPolicyKey;
    }
 
    template<typename ...ComponentTypes>
    FEntityGroupingPolicyKey AddGrouping(TComponentTypeID<ComponentTypes>... InComponents)
    {
        UE::MovieScene::TTupleGroupingPolicy<ComponentTypes...> TuplePolicy;
        return AddGrouping(MoveTemp(TuplePolicy), InComponents...);
    }
 
    UE_API void RemoveGrouping(FEntityGroupingPolicyKey InPolicyKey);
 
    UE_API int32 AllocateGroupIndex(FEntityGroupingPolicyKey InPolicy);
 
 
    int32 NumGroups() const
    {
        return AllocatedGroupIndices.GetMaxIndex();
    }
 
    void FreeEmptyGroups();
 
private:
 
    virtual bool IsRelevantImpl(UMovieSceneEntitySystemLinker* InLinker) const override;
    virtual void OnRun(FSystemTaskPrerequisites& InPrerequisites, FSystemSubsequentTasks& Subsequents) override;
    virtual void OnLink() override;
    virtual void OnUnlink() override;
    virtual void OnCleanTaggedGarbage() override;
 
    void ProcessModifiedGroups();
 
#if WITH_EDITOR
    void OnObjectsReplaced(const TMap<UObject*, UObject*>& ReplacementMap);
#endif
 
private:
 
 
    TSparseArray<FEntityGroupingPolicyKey> AllocatedGroupIndices;
 
    struct FEntityGroupInfo
    {
        int32 NumEntities = 0;
    };
 
    TMap<FEntityGroupID, FEntityGroupInfo> Groups;
 
    TSparseArray<FEntityGroupID> EntityIDToGroup;
 
    TBitArray<> EmptyGroupIndices;
 
    struct FEntityGroupingHandlerInfo
    {
        TInlineValue<IEntityGroupingHandler, 256> Handler;
        UE::MovieScene::FEntityComponentFilter ComponentFilter;
    };
 
    TSparseArray<FEntityGroupingHandlerInfo> GroupHandlers;
 
    friend struct UE::MovieScene::FAddGroupMutation;
    friend struct UE::MovieScene::FUpdateGroupsTask;
    friend struct UE::MovieScene::FEntityGroupBuilder;
};
 
 
#undef UE_API