Fixed a bug with type ids for inherited classes with a custom serializer.

The code that adds a $type field for pointers where needed was still assuming that custom serializer were always for primitives, which isn't the case anymore. This changes updates the behavior to allow $type to be added to those as well as long as they use an object. This does now however rely more heavily on earlier checks that the data needs a $type because it otherwise can't tell the difference between a primitive getting a default value (an empty object). In the original code this situation would have resulted in failed serialization though, so it's unlikely to be a problem.

Code/Framework/AzCore/AzCore/Serialization/Json/JsonSerializer.cpp

@@ -104,13 +104,12 @@ namespace AZ
         auto serializer = context.GetRegistrationContext()->GetSerializerForType(classData.m_typeId);
         if (serializer)
         {
-            if (storeTypeId == StoreTypeId::Yes)
+            ResultCode result = serializer->Store(node, object, defaultObject, classData.m_typeId, context);
+            if (storeTypeId == StoreTypeId::Yes && result.GetProcessing() != Processing::Halted)
             {
-                return context.Report(Tasks::WriteValue, Outcomes::Catastrophic,
-                    "Unable to store type information in a JSON Serializer primitive.");
+                result.Combine(InsertTypeId(node, classData, context));
             }
-
-            return serializer->Store(node, object, defaultObject, classData.m_typeId, context);
+            return result;
         }
 
         if (classData.m_azRtti && (classData.m_azRtti->GetTypeTraits() & AZ::TypeTraits::is_enum) == AZ::TypeTraits::is_enum)
@@ -128,13 +127,12 @@ namespace AZ
             serializer = context.GetRegistrationContext()->GetSerializerForType(classData.m_azRtti->GetGenericTypeId());
             if (serializer)
             {
-                if (storeTypeId == StoreTypeId::Yes)
+                ResultCode result = serializer->Store(node, object, defaultObject, classData.m_typeId, context);
+                if (storeTypeId == StoreTypeId::Yes && result.GetProcessing() != Processing::Halted)
                 {
-                    return context.Report(Tasks::WriteValue, Outcomes::Catastrophic,
-                        "Unable to store type information in a JSON Serializer primitive.");
+                    result.Combine(InsertTypeId(node, classData, context));
                 }
-
-                return serializer->Store(node, object, defaultObject, classData.m_typeId, context);
+                return result;
             }
         }
         
@@ -149,6 +147,7 @@ namespace AZ
             ResultCode result(Tasks::WriteValue);
             if (storeTypeId == StoreTypeId::Yes)
             {
+                // Not using InsertTypeId here to avoid needing to create the temporary value and swap it in that call.
                 node.AddMember(rapidjson::StringRef(JsonSerialization::TypeIdFieldIdentifier),
                     StoreTypeName(classData, context), context.GetJsonAllocator());
                 result = ResultCode(Tasks::WriteValue, Outcomes::Success);
@@ -569,6 +568,31 @@ namespace AZ
         }
     }
 
+    JsonSerializationResult::ResultCode JsonSerializer::InsertTypeId(
+        rapidjson::Value& output, const SerializeContext::ClassData& classData, JsonSerializerContext& context)
+    {
+        using namespace JsonSerializationResult;
+
+        if (output.IsObject())
+        {
+            rapidjson::Value insertedObject(rapidjson::kObjectType);
+            insertedObject.AddMember(
+                rapidjson::StringRef(JsonSerialization::TypeIdFieldIdentifier), StoreTypeName(classData, context),
+                context.GetJsonAllocator());
+
+            for (auto& [key, element] : output.GetObject())
+            {
+                insertedObject.AddMember(AZStd::move(key), AZStd::move(element), context.GetJsonAllocator());
+            }
+            output = AZStd::move(insertedObject);
+            return ResultCode(Tasks::WriteValue, Outcomes::Success);
+        }
+        else
+        {
+            return context.Report(Tasks::WriteValue, Outcomes::Catastrophic, "Only able to store type information in a JSON Object.");
+        }
+    }
+
     rapidjson::Value JsonSerializer::GetExplicitDefault()
     {
         return rapidjson::Value(rapidjson::kObjectType);

Code/Framework/AzCore/AzCore/Serialization/Json/JsonSerializer.h

@@ -80,6 +80,9 @@ namespace AZ
         static JsonSerializationResult::ResultCode StoreTypeName(rapidjson::Value& output,
             const Uuid& typeId, JsonSerializerContext& context);
 
+        static JsonSerializationResult::ResultCode InsertTypeId(
+            rapidjson::Value& output, const SerializeContext::ClassData& classData, JsonSerializerContext& context);
+
         static rapidjson::Value GetExplicitDefault();
     };
 } // namespace AZ

Code/Framework/AzCore/Tests/Serialization/Json/JsonSerializationTests.cpp

@@ -449,6 +449,46 @@ namespace JsonSerializationTests
 
     }
 
+    TEST_F(JsonSerializationTests, Load_PrimitiveForInheritedClass_LoadsCorrectClass)
+    {
+        using namespace AZ::JsonSerializationResult;
+        using namespace ::testing;
+
+        AZStd::string json = AZStd::string::format(R"(
+            {
+                "%s": "SimpleInheritence",
+                "base_var": -88.0,
+                "var1": 88,
+                "var2": 42
+            })",
+            AZ::JsonSerialization::TypeIdFieldIdentifier);
+        m_jsonDocument->Parse(json.c_str());
+
+        SimpleInheritence::Reflect(m_serializeContext, true);
+        m_serializeContext->RegisterGenericType<AZStd::unique_ptr<BaseClass>>();
+        m_jsonRegistrationContext->Serializer<JsonSerializerMock>()->HandlesType<SimpleInheritence>();
+        JsonSerializerMock* mock =
+            reinterpret_cast<JsonSerializerMock*>(m_jsonRegistrationContext->GetSerializerForType(azrtti_typeid<SimpleInheritence>()));
+        EXPECT_CALL(*mock, Load(_, _, _, _))
+            .Times(Exactly(1))
+            .WillRepeatedly(Return(Result(m_deserializationSettings->m_reporting, "Test", Tasks::ReadField, Outcomes::Success, "")));
+
+        AZStd::unique_ptr<BaseClass> instance;
+        ResultCode result = AZ::JsonSerialization::Load(instance, *m_jsonDocument, *m_deserializationSettings);
+        EXPECT_NE(Processing::Halted, result.GetProcessing());
+
+        EXPECT_EQ(azrtti_typeid<SimpleInheritence>(), azrtti_typeid(*instance));
+
+        m_serializeContext->EnableRemoveReflection();
+        SimpleInheritence::Reflect(m_serializeContext, true);
+        m_serializeContext->DisableRemoveReflection();
+
+        m_jsonRegistrationContext->EnableRemoveReflection();
+        m_serializeContext->RegisterGenericType<AZStd::unique_ptr<BaseClass>>();
+        m_jsonRegistrationContext->Serializer<JsonSerializerMock>()->HandlesType<SimpleInheritence>();
+        m_jsonRegistrationContext->DisableRemoveReflection();
+    }
+
     TEST_F(JsonSerializationTests, Store_TemplatedClassWithRegisteredHandler_StoreOnHandlerCalled)
     {
         using namespace AZ::JsonSerializationResult;
@@ -474,6 +514,52 @@ namespace JsonSerializationTests
         m_jsonRegistrationContext->DisableRemoveReflection();
     }
 
+    TEST_F(JsonSerializationTests, Store_PrimitiveForInheritedClass_StoreSucceedsAndTypeIdIsAdded)
+    {
+        using namespace AZ::JsonSerializationResult;
+        using namespace ::testing;
+
+        SimpleInheritence::Reflect(m_serializeContext, true);
+        m_serializeContext->RegisterGenericType<AZStd::unique_ptr<BaseClass>>();
+        m_jsonRegistrationContext->Serializer<JsonSerializerMock>()->HandlesType<SimpleInheritence>();
+        JsonSerializerMock* mock =
+            reinterpret_cast<JsonSerializerMock*>(m_jsonRegistrationContext->GetSerializerForType(azrtti_typeid<SimpleInheritence>()));
+        EXPECT_CALL(*mock, Store(_, _, _, _, _))
+            .Times(Exactly(1))
+            .WillRepeatedly(Invoke([](rapidjson::Value& output, const void*, const void*, const AZ::Uuid&, AZ::JsonSerializerContext& context)
+                {
+                    // Insert some values to allow verification later on.
+                    output.SetObject();
+                    output.AddMember("base_var", -88.0, context.GetJsonAllocator());
+                    output.AddMember("var1", 88, context.GetJsonAllocator());
+                    output.AddMember("var2", 42, context.GetJsonAllocator());
+                    return context.Report(Tasks::WriteValue, Outcomes::Success, "");
+                }));
+
+        AZStd::unique_ptr<BaseClass> instance{ aznew SimpleInheritence() };
+        ResultCode result = AZ::JsonSerialization::Store(*m_jsonDocument, m_jsonDocument->GetAllocator(), instance, *m_serializationSettings);
+        EXPECT_NE(Processing::Halted, result.GetProcessing());
+
+        AZStd::string compare = AZStd::string::format(R"(
+            {
+                "%s": "SimpleInheritence",
+                "base_var": -88.0,
+                "var1": 88,
+                "var2": 42
+            })",
+            AZ::JsonSerialization::TypeIdFieldIdentifier);
+        Expect_DocStrEq(compare.c_str());
+
+        m_serializeContext->EnableRemoveReflection();
+        m_serializeContext->RegisterGenericType<AZStd::unique_ptr<BaseClass>>();
+        SimpleInheritence::Reflect(m_serializeContext, true);
+        m_serializeContext->DisableRemoveReflection();
+
+        m_jsonRegistrationContext->EnableRemoveReflection();
+        m_jsonRegistrationContext->Serializer<JsonSerializerMock>()->HandlesType<SimpleInheritence>();
+        m_jsonRegistrationContext->DisableRemoveReflection();
+    }
+
     TEST_F(JsonSerializationTests, Store_StoreWithNullPtr_ReturnsCatastrophic)
     {
         using namespace AZ::JsonSerializationResult;