Support value types in complex JSON shaper (#36557)

Closes #36552

Author: roji

src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.ShaperProcessingExpressionVisitor.ClientMethods.cs

@@ -70,8 +70,11 @@ private static readonly MethodInfo IncludeJsonEntityReferenceMethodInfo
         private static readonly MethodInfo IncludeJsonEntityCollectionMethodInfo
             = typeof(ShaperProcessingExpressionVisitor).GetTypeInfo().GetDeclaredMethod(nameof(IncludeJsonEntityCollection))!;
 
-        private static readonly MethodInfo MaterializeJsonEntityMethodInfo
-            = typeof(ShaperProcessingExpressionVisitor).GetTypeInfo().GetDeclaredMethod(nameof(MaterializeJsonEntity))!;
+        private static readonly MethodInfo MaterializeJsonStructuralTypeMethodInfo
+            = typeof(ShaperProcessingExpressionVisitor).GetTypeInfo().GetDeclaredMethod(nameof(MaterializeJsonStructuralType))!;
+
+        private static readonly MethodInfo MaterializeJsonNullableValueStructuralTypeMethodInfo
+            = typeof(ShaperProcessingExpressionVisitor).GetTypeInfo().GetDeclaredMethod(nameof(MaterializeJsonNullableValueStructuralType))!;
 
         private static readonly MethodInfo MaterializeJsonEntityCollectionMethodInfo
             = typeof(ShaperProcessingExpressionVisitor).GetTypeInfo().GetDeclaredMethod(nameof(MaterializeJsonEntityCollection))!;
@@ -959,20 +962,64 @@ static async Task<RelationalDataReader> InitializeReaderAsync(
         ///     doing so can result in application failures when updating to a new Entity Framework Core release.
         /// </summary>
         [EntityFrameworkInternal]
-        public static TEntity? MaterializeJsonEntity<TEntity>(
+        public static TStructural? MaterializeJsonStructuralType<TStructural>(
             QueryContext queryContext,
             object[]? keyPropertyValues,
             JsonReaderData? jsonReaderData,
             bool nullable,
-            Func<QueryContext, object[]?, JsonReaderData, TEntity> shaper)
-            where TEntity : class
+            Func<QueryContext, object[]?, JsonReaderData, TStructural> shaper)
+        {
+            if (jsonReaderData == null)
+            {
+                return nullable
+                    ? default
+                    : throw new InvalidOperationException(
+                        RelationalStrings.JsonRequiredEntityWithNullJson(typeof(TStructural).Name));
+            }
+
+            var manager = new Utf8JsonReaderManager(jsonReaderData, queryContext.QueryLogger);
+            var tokenType = manager.CurrentReader.TokenType;
+
+            switch (tokenType)
+            {
+                case JsonTokenType.Null:
+                    return nullable
+                        ? default
+                        : throw new InvalidOperationException(
+                            RelationalStrings.JsonRequiredEntityWithNullJson(typeof(TStructural).Name));
+
+                case not JsonTokenType.StartObject:
+                    throw new InvalidOperationException(
+                        CoreStrings.JsonReaderInvalidTokenType(tokenType.ToString()));
+            }
+
+            manager.CaptureState();
+            var result = shaper(queryContext, keyPropertyValues, jsonReaderData);
+
+            return result;
+        }
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     any release. You should only use it directly in your code with extreme caution and knowing that
+        ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+        /// </summary>
+        [EntityFrameworkInternal]
+        public static TStructural? MaterializeJsonNullableValueStructuralType<TStructural>(
+            QueryContext queryContext,
+            object[]? keyPropertyValues,
+            JsonReaderData? jsonReaderData,
+            bool nullable,
+            Func<QueryContext, object[]?, JsonReaderData, TStructural> shaper)
+            where TStructural : struct
         {
             if (jsonReaderData == null)
             {
                 return nullable
                     ? null
                     : throw new InvalidOperationException(
-                        RelationalStrings.JsonRequiredEntityWithNullJson(typeof(TEntity).Name));
+                        RelationalStrings.JsonRequiredEntityWithNullJson(typeof(TStructural).Name));
             }
 
             var manager = new Utf8JsonReaderManager(jsonReaderData, queryContext.QueryLogger);
@@ -984,7 +1031,7 @@ static async Task<RelationalDataReader> InitializeReaderAsync(
                     return nullable
                         ? null
                         : throw new InvalidOperationException(
-                            RelationalStrings.JsonRequiredEntityWithNullJson(typeof(TEntity).Name));
+                            RelationalStrings.JsonRequiredEntityWithNullJson(typeof(TStructural).Name));
 
                 case not JsonTokenType.StartObject:
                     throw new InvalidOperationException(
@@ -1081,16 +1128,14 @@ static async Task<RelationalDataReader> InitializeReaderAsync(
         ///     doing so can result in application failures when updating to a new Entity Framework Core release.
         /// </summary>
         [EntityFrameworkInternal]
-        public static void IncludeJsonEntityReference<TIncludingEntity, TIncludedEntity>(
+        public static void IncludeJsonEntityReference<TStructural, TRelatedStructural>(
             QueryContext queryContext,
             object[]? keyPropertyValues,
             JsonReaderData? jsonReaderData,
-            TIncludingEntity entity,
-            Func<QueryContext, object[]?, JsonReaderData, TIncludedEntity> innerShaper,
-            Action<TIncludingEntity, TIncludedEntity> fixup,
+            TStructural structuralType,
+            Func<QueryContext, object[]?, JsonReaderData, TRelatedStructural> innerShaper,
+            Action<TStructural, TRelatedStructural> fixup,
             bool performFixup)
-            where TIncludingEntity : class
-            where TIncludedEntity : class
         {
             if (jsonReaderData == null)
             {
@@ -1114,7 +1159,7 @@ public static void IncludeJsonEntityReference<TIncludingEntity, TIncludedEntity>
 
             if (performFixup)
             {
-                fixup(entity, included);
+                fixup(structuralType, included);
             }
         }
 

src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.ShaperProcessingExpressionVisitor.cs

@@ -1672,7 +1672,7 @@ private Expression CreateJsonShapers(
                         var elementFixup = Lambda(
                             Block(
                                 typeof(void),
-                                AssignReferenceRelationship(
+                                AssignStructuralProperty(
                                     innerFixupCollectionElementParameter,
                                     innerFixupParentParameter,
                                     inverseNavigation)),
@@ -1709,7 +1709,7 @@ private Expression CreateJsonShapers(
                 {
                     var fixup = GenerateReferenceFixupForJson(
                         structuralType.ClrType,
-                        relatedStructuralType.ClrType,
+                        nestedRelationship.ClrType,
                         nestedRelationship,
                         inverseNavigation);
 
@@ -1847,8 +1847,27 @@ private Expression CreateJsonShapers(
                 return materializeJsonEntityCollectionMethodCall;
             }
 
+
+            // Return the materializer for this JSON object, including null checks which would return null.
+            MethodInfo method;
+
+            if (relationship is not null && Nullable.GetUnderlyingType(relationship.ClrType) is { } underlyingType)
+            {
+                // The association property into which we're assigning has a nullable value type, so generate
+                // a materializer that returns that nullable value type (note that the shaperLambda that
+                // we pass itself always returns a non-nullable value (the null checks are outside of it.))
+                Check.DebugAssert(nullable, "On non-nullable relationship but the relationship's ClrType is Nullable<T>");
+                Check.DebugAssert(underlyingType == structuralType.ClrType);
+
+                method = MaterializeJsonNullableValueStructuralTypeMethodInfo.MakeGenericMethod(structuralType.ClrType);
+            }
+            else
+            {
+                method = MaterializeJsonStructuralTypeMethodInfo.MakeGenericMethod(structuralType.ClrType);
+            }
+
             var materializedRootJsonEntity = Call(
-                MaterializeJsonEntityMethodInfo.MakeGenericMethod(structuralType.ClrType),
+                method,
                 QueryCompilationContext.QueryContextParameter,
                 keyValuesParameter,
                 jsonReaderDataParameter,
@@ -1969,9 +1988,9 @@ protected override Expression VisitSwitch(SwitchExpression switchExpression)
 
                     var managerVariable = Variable(typeof(Utf8JsonReaderManager), "jsonReaderManager");
                     var tokenTypeVariable = Variable(typeof(JsonTokenType), "tokenType");
-                    var jsonEntityTypeVariable = (ParameterExpression)jsonEntityTypeInitializerBlock.Expressions[^1];
+                    var jsonStructuralTypeVariable = (ParameterExpression)jsonEntityTypeInitializerBlock.Expressions[^1];
 
-                    Debug.Assert(jsonEntityTypeVariable.Type == structuralType.ClrType);
+                    Debug.Assert(jsonStructuralTypeVariable.Type == structuralType.ClrType);
 
                     var finalBlockVariables = new List<ParameterExpression>
                     {
@@ -2024,7 +2043,7 @@ protected override Expression VisitSwitch(SwitchExpression switchExpression)
                     // - navigation fixups
                     // - entity instance variable that is returned as end result
                     var propertyAssignmentReplacer = new ValueBufferTryReadValueMethodsReplacer(
-                        jsonEntityTypeVariable, propertyAssignmentMap);
+                        jsonStructuralTypeVariable, propertyAssignmentMap);
 
                     if (body.Expressions[0] is BinaryExpression
                         {
@@ -2051,7 +2070,7 @@ protected override Expression VisitSwitch(SwitchExpression switchExpression)
                     // or for empty/null collections of a tracking queries.
                     ProcessFixup(queryStateManager ? trackingInnerFixupMap : innerFixupMap);
 
-                    finalBlockExpressions.Add(jsonEntityTypeVariable);
+                    finalBlockExpressions.Add(jsonStructuralTypeVariable);
 
                     return Block(
                         finalBlockVariables,
@@ -2063,18 +2082,35 @@ void ProcessFixup(IDictionary<string, LambdaExpression> fixupMap)
                         {
                             var navigationEntityParameter = _navigationVariableMap[fixup.Key];
 
-                            // we need to add null checks before we run fixup logic. For regular entities, whose fixup is done as part of the "Materialize*" method
-                            // the checks are done there (same will be done for the "optimized" scenario, where we populate properties directly rather than store in variables)
-                            // but in this case fixups are standalone, so the null safety must be added by us directly
-                            finalBlockExpressions.Add(
-                                IfThen(
-                                    NotEqual(
-                                        jsonEntityTypeVariable,
-                                        Constant(null, jsonEntityTypeVariable.Type)),
-                                    Invoke(
-                                        fixup.Value,
-                                        jsonEntityTypeVariable,
-                                        _navigationVariableMap[fixup.Key])));
+                            // Inject the fixup code for each property; we have this as a set of lambdas in the fixup map.
+                            // In the normal case, simply Invoke the lambda, passing it the structural type to be fixed up as a parameter.
+                            // This unfortunately doesn't work on value types (where a copy would be mutated), so for them,
+                            // we unwrap the lambda and integrate its body directly.
+                            // We should ideally do this for all cases (no need for the extra lambda Invoke), but there are some issues around us writing
+                            // to readonly fields.
+                            if (jsonStructuralTypeVariable.Type.IsValueType /*&& Nullable.GetUnderlyingType(jsonStructuralTypeVariable.Type) is null*/)
+                            {
+                                var fixupBody = ReplacingExpressionVisitor.Replace(
+                                    originals: [fixup.Value.Parameters[0], fixup.Value.Parameters[1]],
+                                    replacements: [jsonStructuralTypeVariable, _navigationVariableMap[fixup.Key]],
+                                    fixup.Value.Body);
+
+                                finalBlockExpressions.Add(fixupBody);
+                            }
+                            else
+                            {
+                                // If the structural type being fixed up is nullable, then we need to add null checks before we run fixup logic.
+                                // For regular entities, whose fixup is done as part of the "Materialize*" method, the checks are done there
+                                // (the same will be done for the "optimized" scenario, where we populate properties directly rather than store in variables).
+                                // But in this case fixups are standalone, so the null safety must be added here.
+                                finalBlockExpressions.Add(
+                                    IfThen(
+                                        NotEqual(jsonStructuralTypeVariable, Constant(null, jsonStructuralTypeVariable.Type)),
+                                        Invoke(
+                                            fixup.Value,
+                                            jsonStructuralTypeVariable,
+                                            _navigationVariableMap[fixup.Key])));
+                            }
                         }
                     }
                 }
@@ -2756,7 +2792,7 @@ private LambdaExpression GenerateFixup(
                 expressions.Add(
                     relationship.IsCollection
                         ? AddToCollectionRelationship(entityParameter, relatedEntityParameter, relationship)
-                        : AssignReferenceRelationship(entityParameter, relatedEntityParameter, relationship));
+                        : AssignStructuralProperty(entityParameter, relatedEntityParameter, relationship));
             }
 
             if (inverseNavigation != null
@@ -2765,26 +2801,26 @@ private LambdaExpression GenerateFixup(
                 expressions.Add(
                     inverseNavigation.IsCollection
                         ? AddToCollectionRelationship(relatedEntityParameter, entityParameter, inverseNavigation)
-                        : AssignReferenceRelationship(relatedEntityParameter, entityParameter, inverseNavigation));
+                        : AssignStructuralProperty(relatedEntityParameter, entityParameter, inverseNavigation));
             }
 
             return Lambda(Block(typeof(void), expressions), entityParameter, relatedEntityParameter);
         }
 
         private static LambdaExpression GenerateReferenceFixupForJson(
-            Type entityType,
-            Type relatedEntityType,
+            Type clrType,
+            Type relatedClrType,
             IPropertyBase relationship,
             INavigationBase? inverseNavigation)
         {
-            var entityParameter = Parameter(entityType);
-            var relatedEntityParameter = Parameter(relatedEntityType);
+            var entityParameter = Parameter(clrType);
+            var relatedEntityParameter = Parameter(relatedClrType);
             var expressions = new List<Expression>();
 
             if (!relationship.IsShadowProperty())
             {
                 expressions.Add(
-                    AssignReferenceRelationship(
+                    AssignStructuralProperty(
                         entityParameter,
                         relatedEntityParameter,
                         relationship));
@@ -2794,7 +2830,7 @@ private static LambdaExpression GenerateReferenceFixupForJson(
                 && !inverseNavigation.IsShadowProperty())
             {
                 expressions.Add(
-                    AssignReferenceRelationship(
+                    AssignStructuralProperty(
                         relatedEntityParameter,
                         entityParameter,
                         inverseNavigation));
@@ -2821,11 +2857,21 @@ public static void InverseCollectionFixup<TCollectionElement, TEntity>(
             }
         }
 
-        private static Expression AssignReferenceRelationship(
-            ParameterExpression entity,
-            ParameterExpression relatedEntity,
-            IPropertyBase relationship)
-            => entity.MakeMemberAccess(relationship.GetMemberInfo(forMaterialization: true, forSet: true)).Assign(relatedEntity);
+        private static Expression AssignStructuralProperty(
+            ParameterExpression structuralType,
+            ParameterExpression relatedStructuralType,
+            IPropertyBase structuralProperty)
+        {
+            var setter = structuralProperty.GetMemberInfo(forMaterialization: true, forSet: true);
+
+            // If we're assigning a value complex type to a nullable complex property, add an upcast for typing
+            var assignee = structuralProperty.ClrType.IsNullableValueType()
+                && structuralProperty.ClrType.UnwrapNullableType() == relatedStructuralType.Type
+                ? Convert(relatedStructuralType, structuralProperty.ClrType)
+                : (Expression)relatedStructuralType;
+
+            return structuralType.MakeMemberAccess(setter).Assign(assignee);
+        }
 
         private Expression GetOrCreateCollectionObjectLambda(Type entityType, IPropertyBase relationship)
         {

src/EFCore.Relational/Query/SqlExpressions/SelectExpression.cs

@@ -2885,14 +2885,14 @@ public static Expression GenerateComplexPropertyShaperExpression(
                     complexType.GetJsonPropertyName()
                     ?? throw new UnreachableException($"No JSON property name for complex property {complexProperty.Name}"),
                     tableAlias,
-                    complexProperty.ClrType,
+                    complexProperty.ClrType.UnwrapNullableType(),
                     typeMapping: containerColumn.StoreTypeMapping,
                     isComplexTypeNullable)
                 : new ColumnExpression(
                     containerColumn.Name,
                     tableAlias,
                     containerColumn,
-                    complexProperty.ClrType,
+                    complexProperty.ClrType.UnwrapNullableType(),
                     containerColumn.StoreTypeMapping,
                     isComplexTypeNullable);