Fix nested prefab bug in detach prefab operation (#14483)

* Fix nested prefab bug in detach prefab operation

Signed-off-by: Junhao Wang <[email protected]>

* Update the simple test to check hierarchy after detach

Signed-off-by: Junhao Wang <[email protected]>

* Add two new unit tests to check top level entity

Signed-off-by: Junhao Wang <[email protected]>

* Fix unused warning

Signed-off-by: Junhao Wang <[email protected]>

* Use ASSERT_TRUE for level instance check

Signed-off-by: Junhao Wang <[email protected]>

---------

Signed-off-by: Junhao Wang <[email protected]>

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabPublicHandler.cpp

@@ -1586,7 +1586,20 @@ namespace AzToolsFramework
                 AZStd::vector<const AZ::Entity*> detachedEntitiesToUpdate;
                 detachedEntitiesToUpdate.push_back(&containerEntity);
 
-                // Detach nested instances and add them to the parent instance
+                // Detach entities and add them to the parent instance.
+                detachedInstance->DetachEntities(
+                    [&detachedEntitiesToUpdate, &parentInstance](AZStd::unique_ptr<AZ::Entity> entityPtr)
+                    {
+                        AZ::Entity* detachedEntity = entityPtr.release();
+                        [[maybe_unused]] const bool entityAdded = parentInstance.AddEntity(*detachedEntity);
+                        AZ_Assert(entityAdded, "Add target Instance's entity to its parent Instance failed.");
+
+                        detachedEntitiesToUpdate.push_back(detachedEntity);
+                    });
+
+                // Detach nested instances and add them to the parent instance.
+                // This step needs to happen after detaching entities to make sure the parent entity reference inside nested instance
+                // is up to date since its parent entity is also moved to a new owning instance.
                 detachedInstance->DetachNestedInstances(
                     [&](AZStd::unique_ptr<Instance> detachedNestedInstance)
                     {
@@ -1605,20 +1618,12 @@ namespace AzToolsFramework
                         m_instanceToTemplateInterface->GeneratePatch(
                             reparentPatch, nestedInstanceTemplateDom, nestedInstanceDomUnderNewParent);
 
-                        CreateLink(nestedInstanceUnderNewParent, parentTemplateId, undoBatch.GetUndoBatch(), AZStd::move(reparentPatch), true);
-                    });
-
-                // Detach entities and add them to the parent instance, then update entity DOMs in template
-                detachedInstance->DetachEntities(
-                    [&detachedEntitiesToUpdate, &parentInstance](AZStd::unique_ptr<AZ::Entity> entityPtr)
-                    {
-                        AZ::Entity* detachedEntity = entityPtr.release();
-                        [[maybe_unused]] const bool entityAdded = parentInstance.AddEntity(*detachedEntity);
-                        AZ_Assert(entityAdded, "Add target Instance's entity to its parent Instance failed.");
-
-                        detachedEntitiesToUpdate.push_back(detachedEntity);
+                        // Create link and update template with the new instance DOM.
+                        CreateLink(
+                            nestedInstanceUnderNewParent, parentTemplateId, undoBatch.GetUndoBatch(), AZStd::move(reparentPatch), true);
                     });
 
+                // Update template with the new entity DOMs.
                 PrefabUndoHelpers::AddEntityDoms(detachedEntitiesToUpdate, parentInstance.GetTemplateId(), undoBatch.GetUndoBatch());
 
                 // Update parent entity of the container entity in template with the new sort order information

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

@@ -18,15 +18,18 @@ namespace UnitTest
         // Level
         // | Car       (prefab)  <-- detach prefab
         //   | Tire
+        //     | Belt
 
         const AZStd::string carPrefabName = "CarPrefab";
         const AZStd::string tireEntityName = "Tire";
+        const AZStd::string beltEntityName = "Belt";
 
         AZ::IO::Path engineRootPath;
         m_settingsRegistryInterface->Get(engineRootPath.Native(), AZ::SettingsRegistryMergeUtils::FilePathKey_EngineRootFolder);
         AZ::IO::Path carPrefabFilepath = engineRootPath / carPrefabName;
 
         AZ::EntityId tireEntityId = CreateEditorEntityUnderRoot(tireEntityName);
+        CreateEditorEntity(beltEntityName, tireEntityId);
         AZ::EntityId carContainerId = CreateEditorPrefab(carPrefabFilepath, { tireEntityId });
 
         InstanceAlias carInstanceAlias = FindNestedInstanceAliasInInstance(GetRootContainerEntityId(), carPrefabName);
@@ -41,10 +44,10 @@ namespace UnitTest
         ValidateNestedInstanceNotUnderInstance(GetRootContainerEntityId(), carInstanceAlias);
 
         InstanceOptionalReference levelInstance = m_instanceEntityMapperInterface->FindOwningInstance(GetRootContainerEntityId());
-        EXPECT_TRUE(levelInstance.has_value());
+        ASSERT_TRUE(levelInstance.has_value());
 
-        // Validate there are two entities in the level prefab instance.
-        EXPECT_EQ(levelInstance->get().GetEntityAliasCount(), 2);
+        // Validate there are three entities in the level prefab instance.
+        EXPECT_EQ(levelInstance->get().GetEntityAliasCount(), 3);
 
         // Validate that the car's parent entity is the level container entity.
         AZStd::string carEntityAliasAfterDetach = FindEntityAliasInInstance(GetRootContainerEntityId(), carPrefabName);
@@ -63,6 +66,15 @@ namespace UnitTest
         AZ::EntityId parentEntityIdForTire;
         AZ::TransformBus::EventResult(parentEntityIdForTire, tireEntityIdAfterDetach, &AZ::TransformInterface::GetParentId);
         EXPECT_EQ(carEntityIdAfterDetach, parentEntityIdForTire);
+
+        // Validate that the belt's parent entity is the tire.
+        AZStd::string beltEntityAliasAfterDetach = FindEntityAliasInInstance(GetRootContainerEntityId(), beltEntityName);
+        AZ::EntityId beltEntityIdAfterDetach = levelInstance->get().GetEntityId(beltEntityAliasAfterDetach);
+        EXPECT_TRUE(beltEntityIdAfterDetach.IsValid());
+
+        AZ::EntityId parentEntityIdForBelt;
+        AZ::TransformBus::EventResult(parentEntityIdForBelt, beltEntityIdAfterDetach, &AZ::TransformInterface::GetParentId);
+        EXPECT_EQ(tireEntityIdAfterDetach, parentEntityIdForBelt);
     }
 
     TEST_F(PrefabDetachPrefabTests, DetachPrefabUnderParentSucceeds)
@@ -96,7 +108,7 @@ namespace UnitTest
         ValidateNestedInstanceNotUnderInstance(GetRootContainerEntityId(), carInstanceAlias);
 
         InstanceOptionalReference levelInstance = m_instanceEntityMapperInterface->FindOwningInstance(GetRootContainerEntityId());
-        EXPECT_TRUE(levelInstance.has_value());
+        ASSERT_TRUE(levelInstance.has_value());
 
         // Validate there are three entities in the level prefab instance.
         EXPECT_EQ(levelInstance->get().GetEntityAliasCount(), 3);        
@@ -166,7 +178,7 @@ namespace UnitTest
         ValidateNestedInstanceUnderInstance(GetRootContainerEntityId(), wheelInstanceAlias);
 
         InstanceOptionalReference levelInstance = m_instanceEntityMapperInterface->FindOwningInstance(GetRootContainerEntityId());
-        EXPECT_TRUE(levelInstance.has_value());
+        ASSERT_TRUE(levelInstance.has_value());
         
         AZStd::vector<InstanceOptionalReference> nestedInstances;
         levelInstance->get().GetNestedInstances(
@@ -204,6 +216,95 @@ namespace UnitTest
         EXPECT_EQ(wheelContainerIdAfterDetach, parentEntityIdForTire);
     }
 
+    TEST_F(PrefabDetachPrefabTests, DetachPrefabWithNestedPrefabUnderTopLevelEntitySucceeds)
+    {
+        // Level
+        // | Car          (prefab)   <-- detach prefab
+        //   | Wheels                <-- top level entity
+        //     | Wheel    (prefab)
+        //       | Tire
+
+        const AZStd::string carPrefabName = "CarPrefab";
+        const AZStd::string wheelPrefabName = "WheelPrefab";
+
+        const AZStd::string wheelsEntityName = "Wheels";
+        const AZStd::string tireEntityName = "Tire";
+
+        AZ::IO::Path engineRootPath;
+        m_settingsRegistryInterface->Get(engineRootPath.Native(), AZ::SettingsRegistryMergeUtils::FilePathKey_EngineRootFolder);
+        AZ::IO::Path carPrefabFilepath = engineRootPath / carPrefabName;
+        AZ::IO::Path wheelPrefabFilepath = engineRootPath / wheelPrefabName;
+
+        // Create the wheels and tire entities.
+        AZ::EntityId wheelsEntityId = CreateEditorEntityUnderRoot(wheelsEntityName);
+        AZ::EntityId tireEntityId = CreateEditorEntity(tireEntityName, wheelsEntityId);
+
+        // Create the wheel prefab.
+        AZ::EntityId wheelContainerId = CreateEditorPrefab(wheelPrefabFilepath, { tireEntityId });
+
+        EntityAlias tireEntityAlias = FindEntityAliasInInstance(wheelContainerId, tireEntityName);
+
+        // Create the car prefab.
+        AZ::EntityId carContainerId = CreateEditorPrefab(carPrefabFilepath, { wheelsEntityId });
+        InstanceAlias carInstanceAlias = FindNestedInstanceAliasInInstance(GetRootContainerEntityId(), carPrefabName);
+
+        // Detach the car prefab.
+        PrefabOperationResult result = m_prefabPublicInterface->DetachPrefab(carContainerId);
+        ASSERT_TRUE(result.IsSuccess());
+
+        PropagateAllTemplateChanges();
+
+        // Validate there is no car instance in the level prefab instance.
+        ValidateNestedInstanceNotUnderInstance(GetRootContainerEntityId(), carInstanceAlias);
+
+        // Validate there is wheels entity in the level prefab instance.
+        EntityAlias wheelsEntityAlias = FindEntityAliasInInstance(GetRootContainerEntityId(), wheelsEntityName);
+        ValidateEntityUnderInstance(GetRootContainerEntityId(), wheelsEntityAlias, wheelsEntityName);
+
+        // Validate there is wheel instance in the level prefab instance.
+        InstanceAlias wheelInstanceAlias = FindNestedInstanceAliasInInstance(GetRootContainerEntityId(), wheelPrefabName);
+        ValidateNestedInstanceUnderInstance(GetRootContainerEntityId(), wheelInstanceAlias);
+
+        InstanceOptionalReference levelInstance = m_instanceEntityMapperInterface->FindOwningInstance(GetRootContainerEntityId());
+        ASSERT_TRUE(levelInstance.has_value());
+
+        AZStd::vector<InstanceOptionalReference> nestedInstances;
+        levelInstance->get().GetNestedInstances(
+            [&nestedInstances](AZStd::unique_ptr<Instance>& nestedInstance)
+            {
+                nestedInstances.push_back(*(nestedInstance.get()));
+            });
+
+        EXPECT_EQ(nestedInstances.size(), 1) << "There should be only one nested instance in level after detaching.";
+        EXPECT_TRUE(nestedInstances[0].has_value());
+
+        // Validate that the car's parent entity is the level container entity.
+        AZStd::string carEntityAliasAfterDetach = FindEntityAliasInInstance(GetRootContainerEntityId(), carPrefabName);
+        AZ::EntityId carEntityIdAfterDetach = levelInstance->get().GetEntityId(carEntityAliasAfterDetach);
+        EXPECT_TRUE(carEntityIdAfterDetach.IsValid());
+
+        AZ::EntityId parentEntityIdForCar;
+        AZ::TransformBus::EventResult(parentEntityIdForCar, carEntityIdAfterDetach, &AZ::TransformInterface::GetParentId);
+        EXPECT_EQ(levelInstance->get().GetContainerEntityId(), parentEntityIdForCar);
+
+        // Validate that the wheels' parent entity is the car.
+        AZ::EntityId wheelsEntityIdAfterDetach = levelInstance->get().GetEntityId(wheelsEntityAlias);
+        EXPECT_TRUE(wheelsEntityIdAfterDetach.IsValid());
+
+        AZ::EntityId parentEntityIdForWheels;
+        AZ::TransformBus::EventResult(parentEntityIdForWheels, wheelsEntityIdAfterDetach, &AZ::TransformInterface::GetParentId);
+        EXPECT_EQ(carEntityIdAfterDetach, parentEntityIdForWheels);
+
+        // Validate that the wheel prefab's parent entity is the wheels.
+        Instance& wheelInstanceAfterDetach = nestedInstances[0]->get();
+        AZ::EntityId wheelContainerIdAfterDetach = wheelInstanceAfterDetach.GetContainerEntityId();
+        EXPECT_TRUE(wheelContainerIdAfterDetach.IsValid());
+
+        AZ::EntityId parentEntityIdForWheel;
+        AZ::TransformBus::EventResult(parentEntityIdForWheel, wheelContainerIdAfterDetach, &AZ::TransformInterface::GetParentId);
+        EXPECT_EQ(wheelsEntityIdAfterDetach, parentEntityIdForWheel);
+    }
+
     TEST_F(PrefabDetachPrefabTests, DetachPrefabValidatesDetachedContainerEntityOrder)
     {
         // Validate the detached container entity's sort order in its parent.
@@ -320,7 +421,7 @@ namespace UnitTest
         PropagateAllTemplateChanges();
 
         InstanceOptionalReference levelInstance = m_instanceEntityMapperInterface->FindOwningInstance(GetRootContainerEntityId());
-        EXPECT_TRUE(levelInstance.has_value());
+        ASSERT_TRUE(levelInstance.has_value());
 
         // Validate child entity order under the car after detaching the car prefab.
         EntityAlias carEntityAliasAfterDetach = FindEntityAliasInInstance(levelInstance->get().GetContainerEntityId(), carPrefabName);
@@ -341,4 +442,90 @@ namespace UnitTest
             childEntityName, &AZ::ComponentApplicationRequests::GetEntityName, entityOrderArrayAfterDetach[2]);
         EXPECT_EQ(childEntityName, batteryEntityName);
     }
+
+    TEST_F(PrefabDetachPrefabTests, DetachPrefabValidatesTopLevelChildEntityOrder)
+    {
+        // Validate the sort order of child entities and prefabs that are under the top level entity.
+        // 
+        // Level
+        // | Car          (prefab)   <-- detach prefab
+        //   | Wheels                <-- top level entity
+        //     | Red_Wheel
+        //     | Wheel    (prefab)
+        //       | Tire
+        //     | Black_Wheel
+
+        const AZStd::string carPrefabName = "CarPrefab";
+        const AZStd::string wheelPrefabName = "WheelPrefab";
+
+        const AZStd::string wheelsEntityName = "Wheels";
+        const AZStd::string redWheelEntityName = "Red_Wheel";
+        const AZStd::string blackWheelEntityName = "Black_Wheel";
+        const AZStd::string tireEntityName = "Tire";
+
+        AZ::IO::Path engineRootPath;
+        m_settingsRegistryInterface->Get(engineRootPath.Native(), AZ::SettingsRegistryMergeUtils::FilePathKey_EngineRootFolder);
+        AZ::IO::Path carPrefabFilepath = engineRootPath / carPrefabName;
+        AZ::IO::Path wheelPrefabFilepath = engineRootPath / wheelPrefabName;
+
+        // Create the wheels, red wheel and tire entities.
+        AZ::EntityId wheelsEntityId = CreateEditorEntityUnderRoot(wheelsEntityName);
+        CreateEditorEntity(redWheelEntityName, wheelsEntityId);
+        AZ::EntityId tireEntityId = CreateEditorEntity(tireEntityName, wheelsEntityId);
+
+        // Create the wheel prefab.
+        CreateEditorPrefab(wheelPrefabFilepath, { tireEntityId });
+
+        // Create the black wheel entity.
+        CreateEditorEntity(blackWheelEntityName, wheelsEntityId);
+
+        // Create the car prefab.
+        AZ::EntityId carContainerId = CreateEditorPrefab(carPrefabFilepath, { wheelsEntityId });
+
+        InstanceOptionalReference levelInstance = m_instanceEntityMapperInterface->FindOwningInstance(GetRootContainerEntityId());
+        ASSERT_TRUE(levelInstance.has_value());
+
+        // Validate child entity order under wheels before detaching the car prefab.
+        EntityAlias wheelsEntityAlias = FindEntityAliasInInstance(carContainerId, wheelsEntityName);
+        InstanceOptionalReference carInstance = m_instanceEntityMapperInterface->FindOwningInstance(carContainerId);
+        EXPECT_TRUE(carInstance.has_value());
+        wheelsEntityId = carInstance->get().GetEntityId(wheelsEntityAlias);
+
+        AzToolsFramework::EntityOrderArray entityOrderArrayBeforeDetach = AzToolsFramework::GetEntityChildOrder(wheelsEntityId);
+        EXPECT_EQ(entityOrderArrayBeforeDetach.size(), 3);
+
+        AZStd::string childEntityName;
+        AZ::ComponentApplicationBus::BroadcastResult(
+            childEntityName, &AZ::ComponentApplicationRequests::GetEntityName, entityOrderArrayBeforeDetach[0]);
+        EXPECT_EQ(childEntityName, redWheelEntityName);
+        AZ::ComponentApplicationBus::BroadcastResult(
+            childEntityName, &AZ::ComponentApplicationRequests::GetEntityName, entityOrderArrayBeforeDetach[1]);
+        EXPECT_EQ(childEntityName, wheelPrefabName);
+        AZ::ComponentApplicationBus::BroadcastResult(
+            childEntityName, &AZ::ComponentApplicationRequests::GetEntityName, entityOrderArrayBeforeDetach[2]);
+        EXPECT_EQ(childEntityName, blackWheelEntityName);
+
+        // Detach the car prefab.
+        PrefabOperationResult result = m_prefabPublicInterface->DetachPrefab(carContainerId);
+        ASSERT_TRUE(result.IsSuccess());
+
+        PropagateAllTemplateChanges();
+
+        // Validate child entity order under wheels after detaching the car prefab.
+        wheelsEntityAlias = FindEntityAliasInInstance(levelInstance->get().GetContainerEntityId(), wheelsEntityName);
+        wheelsEntityId = levelInstance->get().GetEntityId(wheelsEntityAlias);
+
+        AzToolsFramework::EntityOrderArray entityOrderArrayAfterDetach = AzToolsFramework::GetEntityChildOrder(wheelsEntityId);
+        EXPECT_EQ(entityOrderArrayAfterDetach.size(), 3);
+
+        AZ::ComponentApplicationBus::BroadcastResult(
+            childEntityName, &AZ::ComponentApplicationRequests::GetEntityName, entityOrderArrayAfterDetach[0]);
+        EXPECT_EQ(childEntityName, redWheelEntityName);
+        AZ::ComponentApplicationBus::BroadcastResult(
+            childEntityName, &AZ::ComponentApplicationRequests::GetEntityName, entityOrderArrayAfterDetach[1]);
+        EXPECT_EQ(childEntityName, wheelPrefabName);
+        AZ::ComponentApplicationBus::BroadcastResult(
+            childEntityName, &AZ::ComponentApplicationRequests::GetEntityName, entityOrderArrayAfterDetach[2]);
+        EXPECT_EQ(childEntityName, blackWheelEntityName);
+    }
 } // namespace UnitTest