From e633b5bdc70ca5b4ea3bb8cc7cd55e7e0acd3e8d Mon Sep 17 00:00:00 2001
From: mnaumov <mnaumov@amazon.com>
Date: Fri, 13 Aug 2021 19:43:38 -0700
Subject: [PATCH] Addressing Ronald's feedback

Signed-off-by: mnaumov <mnaumov@amazon.com>
---
 .../Prefab/Instance/InstanceSerializer.cpp           |  3 +++
 .../AzToolsFramework/Prefab/PrefabDomUtils.cpp       |  6 ++++++
 .../AzToolsFramework/Prefab/PrefabDomUtils.h         | 12 +++++++++++-
 .../AzToolsFramework/Prefab/PrefabLoader.cpp         |  2 +-
 4 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/Instance/InstanceSerializer.cpp b/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/Instance/InstanceSerializer.cpp
index 89e0136b12..af42919637 100644
--- a/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/Instance/InstanceSerializer.cpp
+++ b/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/Instance/InstanceSerializer.cpp
@@ -15,6 +15,7 @@
 #include <AzToolsFramework/Prefab/Instance/InstanceEntityMapperInterface.h>
 #include <AzToolsFramework/Prefab/PrefabLoaderInterface.h>
 #include <AzToolsFramework/Prefab/PrefabSystemComponentInterface.h>
+#include <Prefab/PrefabDomUtils.h>
 
 namespace AzToolsFramework
 {
@@ -81,6 +82,8 @@ namespace AzToolsFramework
                 result.Combine(resultInstances);
             }
 
+            PrefabDomUtils::LinkIdMetadata** linkIdMetadata = context.GetMetadata().Find<PrefabDomUtils::LinkIdMetadata*>();
+            if (linkIdMetadata && *linkIdMetadata)
             {
                 AZ::ScopedContextPath subPathSource(context, "m_linkId");
 
diff --git a/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabDomUtils.cpp b/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabDomUtils.cpp
index 86983d7912..c58cb21295 100644
--- a/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabDomUtils.cpp
+++ b/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabDomUtils.cpp
@@ -64,6 +64,12 @@ namespace AzToolsFramework
                     settings.m_keepDefaults = true;
                 }
 
+                if ((flags & StoreInstanceFlags::StoreLinkIds) != StoreInstanceFlags::None)
+                {
+                    LinkIdMetadata linkIdMetadata;
+                    settings.m_metadata.Add(&linkIdMetadata);
+                }
+
                 AZ::JsonSerializationResult::ResultCode result =
                     AZ::JsonSerialization::Store(prefabDom, prefabDom.GetAllocator(), instance, settings);
 
diff --git a/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabDomUtils.h b/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabDomUtils.h
index e773b581dd..fe38eb9e0f 100644
--- a/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabDomUtils.h
+++ b/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabDomUtils.h
@@ -45,7 +45,11 @@ namespace AzToolsFramework
 
                 //! By default an instance will be stored with default values. In cases where we want to store less json without defaults
                 //! such as saving to disk, this flag will control that behavior.
-                StripDefaultValues = 1 << 0
+                StripDefaultValues = 1 << 0,
+
+                //! We do not save linkIds to file. However when loading a level we want to temporarily save
+                //! linkIds to instance dom so any nested prefabs will have linkIds correctly set.
+                StoreLinkIds = 1 << 1
             };
             AZ_DEFINE_ENUM_BITWISE_OPERATORS(StoreInstanceFlags);
 
@@ -138,6 +142,12 @@ namespace AzToolsFramework
                 [[maybe_unused]] const AZStd::string_view printMessage,
                 [[maybe_unused]] const AzToolsFramework::Prefab::PrefabDomValue& prefabDomValue);
 
+            //! An empty struct for passing to JsonSerializerSettings.m_metadata that is consumed by InstanceSerializer::Store.
+            //! If present in metadata, linkIds will be stored to instance dom.
+            struct LinkIdMetadata
+            {
+                AZ_RTTI(LinkIdMetadata, "{8FF7D299-14E3-41D4-90C5-393A240FAE7C}");
+            };
         } // namespace PrefabDomUtils
     } // namespace Prefab
 } // namespace AzToolsFramework
diff --git a/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabLoader.cpp b/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabLoader.cpp
index ea54c9d10c..2b16134744 100644
--- a/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabLoader.cpp
+++ b/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabLoader.cpp
@@ -300,7 +300,7 @@ namespace AzToolsFramework
             }
 
             PrefabDom storedPrefabDom(&loadedTemplateDom->get().GetAllocator());
-            if (!PrefabDomUtils::StoreInstanceInPrefabDom(loadedPrefabInstance, storedPrefabDom))
+            if (!PrefabDomUtils::StoreInstanceInPrefabDom(loadedPrefabInstance, storedPrefabDom, PrefabDomUtils::StoreInstanceFlags::StoreLinkIds))
             {
                 return false;
             }