Fix undo for nested prefab creation by providing link patches to the undo node

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

@@ -96,9 +96,12 @@ namespace AzToolsFramework
                 // target templates of the other instances.
                 for (auto& nestedInstance : instances)
                 {
-                    PrefabUndoHelpers::RemoveLink(
-                        nestedInstance->GetTemplateId(), commonRootEntityOwningInstance->get().GetTemplateId(),
-                        nestedInstance->GetInstanceAlias(), nestedInstance->GetLinkId(), undoBatch.GetUndoBatch());
+                    PrefabOperationResult removeLinkResult = RemoveLink(
+                        nestedInstance, commonRootEntityOwningInstance->get().GetTemplateId(), undoBatch.GetUndoBatch());
+                    if (!removeLinkResult.IsSuccess())
+                    {
+                        return removeLinkResult;
+                    }
                 }
 
                 PrefabUndoHelpers::UpdatePrefabInstance(
@@ -287,6 +290,49 @@ namespace AzToolsFramework
             m_prefabUndoCache.Store(containerEntityId, AZStd::move(containerEntityDomAfter));
         }
 
+        PrefabOperationResult PrefabPublicHandler::RemoveLink(
+            AZStd::unique_ptr<Instance>& sourceInstance, TemplateId targetTemplateId, UndoSystem::URSequencePoint* undoBatch)
+        {
+            LinkReference nestedInstanceLink = m_prefabSystemComponentInterface->FindLink(sourceInstance->GetLinkId());
+            if (!nestedInstanceLink.has_value())
+            {
+                AZ_Assert(false, "A valid link was not found for one of the instances provided as input for the CreatePrefab operation.");
+                return AZ::Failure(
+                    AZStd::string("A valid link was not found for one of the instances provided as input for the CreatePrefab operation."));
+            }
+
+            PrefabDomReference nestedInstanceLinkDom = nestedInstanceLink->get().GetLinkDom();
+            if (!nestedInstanceLinkDom.has_value())
+            {
+                AZ_Assert(
+                    false,
+                    "A valid DOM was not found for the link corresponding to one of the instances provided as input for the "
+                    "CreatePrefab operation.");
+                return AZ::Failure(AZStd::string("A valid DOM was not found for the link corresponding to one of the instances "
+                                                 "provided as input for the CreatePrefab operation."));
+            }
+
+            PrefabDomValueReference nestedInstanceLinkPatches =
+                PrefabDomUtils::FindPrefabDomValue(nestedInstanceLinkDom->get(), PrefabDomUtils::PatchesName);
+            if (!nestedInstanceLinkPatches.has_value())
+            {
+                AZ_Assert(
+                    false,
+                    "A valid DOM for patches was not found for the link corresponding to one of the instances provided as input for the "
+                    "CreatePrefab operation.");
+                return AZ::Failure(AZStd::string("A valid DOM for patcheswas not found for the link corresponding to one of the instances "
+                                                 "provided as input for the CreatePrefab operation."));
+            }
+
+            PrefabDom patchesCopyForUndoSupport;
+            patchesCopyForUndoSupport.CopyFrom(nestedInstanceLinkPatches->get(), patchesCopyForUndoSupport.GetAllocator());
+            PrefabUndoHelpers::RemoveLink(
+                sourceInstance->GetTemplateId(), targetTemplateId, sourceInstance->GetInstanceAlias(), sourceInstance->GetLinkId(),
+                patchesCopyForUndoSupport, undoBatch);
+
+            return AZ::Success();
+        }
+
         PrefabOperationResult PrefabPublicHandler::SavePrefab(AZ::IO::Path filePath)
         {
             auto templateId = m_prefabSystemComponentInterface->GetTemplateIdFromFilePath(filePath.c_str());

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabPublicHandler.h

@@ -82,6 +82,17 @@ namespace AzToolsFramework
                 const EntityList& topLevelEntities, Instance& sourceInstance, TemplateId targetTemplateId,
                 UndoSystem::URSequencePoint* undoBatch, AZ::EntityId commonRootEntityId);
 
+            /**
+             * Removes the link between template of the sourceInstance and the template corresponding to targetTemplateId.
+             *
+             * \param sourceInstance The instance corresponding to the source template of the link to be removed.
+             * \param targetTemplateId The id of the target template of the link to be removed.
+             * \param undoBatch The undo batch to set as parent for this remove link action.
+             * \return PrefabOperationResult Indicates whether the removal of link was successful.
+             */
+            PrefabOperationResult RemoveLink(
+                AZStd::unique_ptr<Instance>& sourceInstance, TemplateId targetTemplateId, UndoSystem::URSequencePoint* undoBatch);
+
             /**
              * Given a list of entityIds, finds the prefab instance that owns the common root entity of the entityIds.
              * 

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabSystemComponent.cpp

@@ -663,8 +663,9 @@ namespace AzToolsFramework
             newLink.SetSourceTemplateId(linkSourceId);
             newLink.SetInstanceName(instanceAlias.c_str());
             newLink.GetLinkDom().SetObject();
-            newLink.GetLinkDom().AddMember(rapidjson::StringRef(PrefabDomUtils::SourceName),
-                rapidjson::StringRef(sourceTemplate.GetFilePath().c_str()), newLink.GetLinkDom().GetAllocator());
+            newLink.GetLinkDom().AddMember(
+                rapidjson::StringRef(PrefabDomUtils::SourceName), rapidjson::StringRef(sourceTemplate.GetFilePath().c_str()),
+                newLink.GetLinkDom().GetAllocator());
 
             if (linkPatch && linkPatch->get().IsArray() && !(linkPatch->get().Empty()))
             {

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabUndo.cpp

@@ -113,7 +113,7 @@ namespace AzToolsFramework
             , m_sourceId(InvalidTemplateId)
             , m_instanceAlias("")
             , m_linkId(InvalidLinkId)
-            , m_linkDom(PrefabDom())
+            , m_linkPatches(PrefabDom())
             , m_linkStatus(LinkStatus::LINKSTATUS)
         {
             m_prefabSystemComponentInterface = AZ::Interface<PrefabSystemComponentInterface>::Get();
@@ -124,7 +124,7 @@ namespace AzToolsFramework
             const TemplateId& targetId,
             const TemplateId& sourceId,
             const InstanceAlias& instanceAlias,
-            PrefabDomReference linkDom,
+            PrefabDomReference linkPatches,
             const LinkId linkId)
         {
             m_targetId = targetId;
@@ -132,9 +132,9 @@ namespace AzToolsFramework
             m_instanceAlias = instanceAlias;
             m_linkId = linkId;
 
-            if (linkDom.has_value())
+            if (linkPatches.has_value())
             {
-                m_linkDom = AZStd::move(linkDom->get());
+                m_linkPatches = AZStd::move(linkPatches->get());
             }
 
             //if linkId is invalid, set as ADD
@@ -193,7 +193,7 @@ namespace AzToolsFramework
 
         void PrefabUndoInstanceLink::AddLink()
         {
-            m_linkId = m_prefabSystemComponentInterface->CreateLink(m_targetId, m_sourceId, m_instanceAlias, m_linkDom, m_linkId);
+            m_linkId = m_prefabSystemComponentInterface->CreateLink(m_targetId, m_sourceId, m_instanceAlias, m_linkPatches, m_linkId);
         }
 
         void PrefabUndoInstanceLink::RemoveLink()

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabUndo.h

@@ -101,7 +101,7 @@ namespace AzToolsFramework
                 const TemplateId& targetId,
                 const TemplateId& sourceId,
                 const InstanceAlias& instanceAlias,
-                PrefabDomReference linkDom = PrefabDomReference(),
+                PrefabDomReference linkPatches = PrefabDomReference(),
                 const LinkId linkId = InvalidLinkId);
 
             void Undo() override;
@@ -120,7 +120,7 @@ namespace AzToolsFramework
             InstanceAlias m_instanceAlias;
 
             LinkId m_linkId;
-            PrefabDom m_linkDom;  //data for delete/update
+            PrefabDom m_linkPatches;  //data for delete/update
             LinkStatus m_linkStatus;
 
             PrefabSystemComponentInterface* m_prefabSystemComponentInterface = nullptr;

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabUndoHelpers.cpp

@@ -46,13 +46,11 @@ namespace AzToolsFramework
             }
 
             void RemoveLink(
-                TemplateId sourceTemplateId, TemplateId targetTemplateId, const InstanceAlias& instanceAlias,
-                LinkId linkId, UndoSystem::URSequencePoint* undoBatch)
+                TemplateId sourceTemplateId, TemplateId targetTemplateId, const InstanceAlias& instanceAlias, LinkId linkId,
+                PrefabDomReference linkPatches, UndoSystem::URSequencePoint* undoBatch)
             {
                 auto linkRemoveUndo = aznew PrefabUndoInstanceLink("Remove Link");
-                PrefabDom emptyLinkDom;
-                linkRemoveUndo->Capture(
-                    targetTemplateId, sourceTemplateId, instanceAlias, emptyLinkDom, linkId);
+                linkRemoveUndo->Capture(targetTemplateId, sourceTemplateId, instanceAlias, linkPatches, linkId);
                 linkRemoveUndo->SetParent(undoBatch);
                 linkRemoveUndo->Redo();
             }

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabUndoHelpers.h

@@ -25,8 +25,8 @@ namespace AzToolsFramework
                 TemplateId sourceTemplateId, TemplateId targetTemplateId, PrefabDomReference patch,
                 const InstanceAlias& instanceAlias, UndoSystem::URSequencePoint* undoBatch);
             void RemoveLink(
-                TemplateId sourceTemplateId, TemplateId targetTemplateId, const InstanceAlias& instanceAlias,
-                LinkId linkId, UndoSystem::URSequencePoint* undoBatch);
+                TemplateId sourceTemplateId, TemplateId targetTemplateId, const InstanceAlias& instanceAlias, LinkId linkId,
+                PrefabDomReference linkPatches, UndoSystem::URSequencePoint* undoBatch);
         }
     } // namespace Prefab
 } // namespace AzToolsFramework