Fix nested container entities not getting editor info removed during prefab processing

Code/Framework/AzToolsFramework/AzToolsFramework/Entity/PrefabEditorEntityOwnershipService.cpp

@@ -159,12 +159,23 @@ namespace AzToolsFramework
 
     void PrefabEditorEntityOwnershipService::GetNonPrefabEntities(EntityList& entities)
     {
-        m_rootInstance->GetEntities(entities, false);
+        m_rootInstance->GetEntities(
+            [&entities](const AZStd::unique_ptr<AZ::Entity>& entity)
+            {
+                entities.emplace_back(entity.get());
+                return true;
+            });
     }
 
     bool PrefabEditorEntityOwnershipService::GetAllEntities(EntityList& entities)
     {
-        m_rootInstance->GetEntities(entities, true);
+        m_rootInstance->GetAllEntitiesInHierarchy(
+            [&entities](const AZStd::unique_ptr<AZ::Entity>& entity)
+            {
+                entities.emplace_back(entity.get());
+                return true;
+            });
+
         return true;
     }
 
@@ -544,7 +555,7 @@ namespace AzToolsFramework
                     return;
                 }
 
-                m_rootInstance->GetNestedEntities([this](AZStd::unique_ptr<AZ::Entity>& entity)
+                m_rootInstance->GetAllEntitiesInHierarchy([this](AZStd::unique_ptr<AZ::Entity>& entity)
                     {
                         AZ_Assert(entity, "Invalid entity found in root instance while starting play in editor.");
                         if (entity->GetState() == AZ::Entity::State::Active)

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/Instance/Instance.cpp

@@ -373,17 +373,25 @@ namespace AzToolsFramework
             }
         }
 
-        void Instance::GetConstNestedEntities(const AZStd::function<bool(const AZ::Entity&)>& callback)
+        bool Instance::GetEntities_Impl(const AZStd::function<bool(AZStd::unique_ptr<AZ::Entity>&)>& callback)
         {
-            GetConstEntities(callback);
+            for (auto& [entityAlias, entity] : m_entities)
-
-            for (const auto& [instanceAlias, instance] : m_nestedInstances)
             {
-                instance->GetConstNestedEntities(callback);
+                if (!entity)
+                {
+                    continue;
+                }
+
+                if (!callback(entity))
+                {
+                    return false;
+                }
             }
+
+            return true;
         }
 
-        void Instance::GetConstEntities(const AZStd::function<bool(const AZ::Entity&)>& callback)
+        bool Instance::GetConstEntities_Impl(const AZStd::function<bool(const AZ::Entity&)>& callback) const
         {
             for (const auto& [entityAlias, entity] : m_entities)
             {
@@ -394,19 +402,83 @@ namespace AzToolsFramework
 
                 if (!callback(*entity))
                 {
-                    break;
+                    return false;
                 }
             }
+
+            return true;
         }
 
-        void Instance::GetNestedEntities(const AZStd::function<bool(AZStd::unique_ptr<AZ::Entity>&)>& callback)
+        bool Instance::GetAllEntitiesInHierarchy_Impl(const AZStd::function<bool(AZStd::unique_ptr<AZ::Entity>&)>& callback)
         {
-            GetEntities(callback);
+            if (HasContainerEntity())
+            {
+                if (!callback(m_containerEntity))
+                {
+                    return false;
+                }
+            }
+
+            if (!GetEntities_Impl(callback))
+            {
+                return false;
+            }
 
             for (auto& [instanceAlias, instance] : m_nestedInstances)
             {
-                instance->GetNestedEntities(callback);
+                if (!instance->GetAllEntitiesInHierarchy_Impl(callback))
+                {
+                    return false;
+                }
             }
+
+            return true;
+        }
+
+        bool Instance::GetAllEntitiesInHierarchyConst_Impl(const AZStd::function<bool(const AZ::Entity&)>& callback) const
+        {
+            if (HasContainerEntity())
+            {
+                if (!callback(*m_containerEntity))
+                {
+                    return false;
+                }
+            }
+
+            if (!GetConstEntities_Impl(callback))
+            {
+                return false;
+            }
+
+            for (const auto& [instanceAlias, instance] : m_nestedInstances)
+            {
+                if (!instance->GetAllEntitiesInHierarchyConst_Impl(callback))
+                {
+                    return false;
+                }
+            }
+
+            return true;
+        }
+
+        void Instance::GetEntities(const AZStd::function<bool(AZStd::unique_ptr<AZ::Entity>&)>& callback)
+        {
+            GetEntities_Impl(callback);
+        }
+
+        void Instance::GetConstEntities(const AZStd::function<bool(const AZ::Entity&)>& callback) const
+        {
+            GetConstEntities_Impl(callback);
+        }
+
+        void Instance::GetAllEntitiesInHierarchy(const AZStd::function<bool(AZStd::unique_ptr<AZ::Entity>&)>& callback)
+        {
+            GetAllEntitiesInHierarchy_Impl(callback);
+        }
+
+        void Instance::GetAllEntitiesInHierarchyConst(const AZStd::function<bool(const AZ::Entity&)>& callback) const
+        {
+            GetAllEntitiesInHierarchyConst_Impl(callback);
         }
 
         void Instance::GetNestedInstances(const AZStd::function<void(AZStd::unique_ptr<Instance>&)>& callback)
@@ -417,44 +489,6 @@ namespace AzToolsFramework
             }
         }
 
-        void Instance::GetEntities(const AZStd::function<bool(AZStd::unique_ptr<AZ::Entity>&)>& callback)
-        {
-            for (auto& [entityAlias, entity] : m_entities)
-            {
-                if (!callback(entity))
-                {
-                    break;
-                }
-            }
-        }
-
-        void Instance::GetEntities(EntityList& entities, bool includeNestedEntities)
-        {
-            // Non-recursive traversal of instances
-            AZStd::vector<Instance*> instancesToTraverse = { this };
-            while (!instancesToTraverse.empty())
-            {
-                Instance* currentInstance = instancesToTraverse.back();
-                instancesToTraverse.pop_back();
-                if (includeNestedEntities)
-                {
-                    instancesToTraverse.reserve(instancesToTraverse.size() + currentInstance->m_nestedInstances.size());
-                    for (const auto& instanceByAlias : currentInstance->m_nestedInstances)
-                    {
-                        instancesToTraverse.push_back(instanceByAlias.second.get());
-                    }
-                }
-
-                // Size increases by 1 for each instance because we have to count the container entity also.
-                entities.reserve(entities.size() + currentInstance->m_entities.size() + 1);
-                entities.push_back(m_containerEntity.get());
-                for (const auto& entityByAlias : currentInstance->m_entities)
-                {
-                    entities.push_back(entityByAlias.second.get());
-                }
-            }
-        }
-
         EntityAliasOptionalReference Instance::GetEntityAlias(const AZ::EntityId& id)
         {
             if (m_instanceToTemplateEntityIdMap.count(id))

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/Instance/Instance.h

@@ -121,10 +121,10 @@ namespace AzToolsFramework
             /**
             * Gets the entities in the Instance DOM.  Can recursively trace all nested instances.
             */
-            void GetConstNestedEntities(const AZStd::function<bool(const AZ::Entity&)>& callback);
-            void GetConstEntities(const AZStd::function<bool(const AZ::Entity&)>& callback);
-            void GetNestedEntities(const AZStd::function<bool(AZStd::unique_ptr<AZ::Entity>&)>& callback);
             void GetEntities(const AZStd::function<bool(AZStd::unique_ptr<AZ::Entity>&)>& callback);
+            void GetConstEntities(const AZStd::function<bool(const AZ::Entity&)>& callback) const;
+            void GetAllEntitiesInHierarchy(const AZStd::function<bool(AZStd::unique_ptr<AZ::Entity>&)>& callback);
+            void GetAllEntitiesInHierarchyConst(const AZStd::function<bool(const AZ::Entity&)>& callback) const;
             void GetNestedInstances(const AZStd::function<void(AZStd::unique_ptr<Instance>&)>& callback);
 
             /**
@@ -184,12 +184,6 @@ namespace AzToolsFramework
 
             static InstanceAlias GenerateInstanceAlias();
 
-        protected:
-            /**
-            * Gets the entities owned by this instance
-            */
-            void GetEntities(EntityList& entities, bool includeNestedEntities = false);
-
         private:
             static constexpr const char s_aliasPathSeparator = '/';
 
@@ -197,6 +191,11 @@ namespace AzToolsFramework
 
             void RemoveEntities(const AZStd::function<bool(const AZStd::unique_ptr<AZ::Entity>&)>& filter);
 
+            bool GetEntities_Impl(const AZStd::function<bool(AZStd::unique_ptr<AZ::Entity>&)>& callback);
+            bool GetConstEntities_Impl(const AZStd::function<bool(const AZ::Entity&)>& callback) const;
+            bool GetAllEntitiesInHierarchy_Impl(const AZStd::function<bool(AZStd::unique_ptr<AZ::Entity>&)>& callback);
+            bool GetAllEntitiesInHierarchyConst_Impl(const AZStd::function<bool(const AZ::Entity&)>& callback) const;
+
             bool RegisterEntity(const AZ::EntityId& entityId, const EntityAlias& entityAlias);
             AZStd::unique_ptr<AZ::Entity> DetachEntity(const EntityAlias& entityAlias);
 

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/Spawnable/EditorInfoRemover.cpp

@@ -62,25 +62,16 @@ namespace AzToolsFramework::Prefab::PrefabConversionUtils
         }
     }
 
-    AZStd::vector<AZ::Entity*> EditorInfoRemover::GetEntitiesFromInstance(AZStd::unique_ptr<Instance>& instance)
+    void EditorInfoRemover::GetEntitiesFromInstance(
+        AZStd::unique_ptr<AzToolsFramework::Prefab::Instance>& instance, EntityList& hierarchyEntities)
     {
-        AZStd::vector<AZ::Entity*> result;
+        instance->GetAllEntitiesInHierarchy(
-
+            [&hierarchyEntities](const AZStd::unique_ptr<AZ::Entity>& entity)
-        instance->GetNestedEntities(
-            [&result](const AZStd::unique_ptr<AZ::Entity>& entity)
             {
-                result.emplace_back(entity.get());
+                hierarchyEntities.emplace_back(entity.get());
                 return true;
             }
         );
-
-        if (instance->HasContainerEntity())
-        {
-            auto containerEntityReference = instance->GetContainerEntity();
-            result.emplace_back(&containerEntityReference->get());
-        }
-
-        return result;
     }
 
     void EditorInfoRemover::SetEditorOnlyEntityHandlerFromCandidates(const EntityList& entities)
@@ -543,7 +534,9 @@ exportComponent, prefabProcessorContext);
         }
 
         // grab all nested entities from the Instance as source entities.
-        EntityList sourceEntities = GetEntitiesFromInstance(instance);
+        EntityList sourceEntities;
+        GetEntitiesFromInstance(instance, sourceEntities);
+
         EntityList exportEntities;
 
         // prepare for validation of component requirements.
@@ -616,7 +609,7 @@ exportComponent, prefabProcessorContext);
         );
 
         // replace entities of instance with exported ones.
-        instance->GetNestedEntities(
+        instance->GetAllEntitiesInHierarchy(
             [&exportEntitiesMap](AZStd::unique_ptr<AZ::Entity>& entity)
             {
                 auto entityId = entity->GetId();
@@ -625,14 +618,6 @@ exportComponent, prefabProcessorContext);
             }
         );
 
-        if (instance->HasContainerEntity())
-        {
-            if (auto found = exportEntitiesMap.find(instance->GetContainerEntityId()); found != exportEntitiesMap.end())
-            {
-                instance->SetContainerEntity(*found->second);
-            }
-        }
-
         // save the final result in the target Prefab DOM.
         PrefabDom filteredPrefab;
         if (!PrefabDomUtils::StoreInstanceInPrefabDom(*instance, filteredPrefab))

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/Spawnable/EditorInfoRemover.h

@@ -55,8 +55,8 @@ namespace AzToolsFramework::Prefab::PrefabConversionUtils
 
      protected:
         using EntityList = AZStd::vector<AZ::Entity*>;
-        static EntityList GetEntitiesFromInstance(
+        static void GetEntitiesFromInstance(
-            AZStd::unique_ptr<AzToolsFramework::Prefab::Instance>& instance);
+            AZStd::unique_ptr<AzToolsFramework::Prefab::Instance>& instance, EntityList& hierarchyEntities);
 
         static bool ReadComponentAttribute(
             AZ::Component* component,

Code/Framework/AzToolsFramework/Tests/Prefab/PrefabUpdateWithPatchesTests.cpp

@@ -93,7 +93,7 @@ namespace UnitTest
 
         // Retrieve the entity pointer from the component application bus.
         AZ::Entity* wheelEntityUnderAxle = nullptr;
-        axleInstance->GetNestedEntities([&wheelEntityUnderAxle, wheelEntityIdUnderAxle](AZStd::unique_ptr<AZ::Entity>& entity)
+        axleInstance->GetAllEntitiesInHierarchy([&wheelEntityUnderAxle, wheelEntityIdUnderAxle](AZStd::unique_ptr<AZ::Entity>& entity)
             {
                 if (entity->GetId() == wheelEntityIdUnderAxle)
                 {