Fix to #31365 - Query: add support for projecting JSON entities that have been composed on (#31391)

Adding new materialization path for JSON entities that have been converted to query roots - they look like normal entity, but materialization is somewhat different.
We need to prune synthesized key properties from materializer code (as they are not associated with any column, but rather made up on the fly)
Also added code to include all the child navigations. For normal entities we have IncludeExpressions in the shaper expression, but for JSON we prune all includes and build them in the materializer.

Fix is to recognize the scenario (entity projection but the entity is mapped to JSON) during apply projection phase and generate all the necessary info needed rather than just dictionary of property to index maps like we do for regular entity.
Then when we build materializer we use that info to prune the synthesized properties from key check and re-use existing include JSON entity code to add child navigations.

Fixes #31365

Author: maumar

src/EFCore.Relational/Query/Internal/QueryableJsonProjectionInfo.cs

@@ -0,0 +1,51 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.EntityFrameworkCore.Query.Internal;
+
+/// <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>
+public readonly struct QueryableJsonProjectionInfo
+{
+    /// <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>
+    public QueryableJsonProjectionInfo(
+        Dictionary<IProperty, int> propertyIndexMap,
+        List<(JsonProjectionInfo, INavigation)> childrenProjectionInfo)
+    {
+        PropertyIndexMap = propertyIndexMap;
+        ChildrenProjectionInfo = childrenProjectionInfo;
+    }
+
+    /// <summary>
+    ///     Map between entity properties and corresponding column indexes.
+    /// </summary>
+    /// <remarks>
+    ///     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.
+    /// </remarks>
+    public IDictionary<IProperty, int> PropertyIndexMap { get; }
+
+    /// <summary>
+    ///     Information needed to construct each child JSON entity.
+    ///     - JsonProjection info (same one we use for simple JSON projection),
+    ///     - navigation between parent and the child JSON entity.
+    /// </summary>
+    /// <remarks>
+    ///     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.
+    /// </remarks>
+    public IList<(JsonProjectionInfo JsonProjectionInfo, INavigation Navigation)> ChildrenProjectionInfo { get; }
+}

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

@@ -433,7 +433,11 @@ protected override Expression VisitBinary(BinaryExpression binaryExpression)
 
                     if (newExpression.Arguments[0] is ProjectionBindingExpression projectionBindingExpression)
                     {
-                        var propertyMap = (IDictionary<IProperty, int>)GetProjectionIndex(projectionBindingExpression);
+                        var projectionIndex = GetProjectionIndex(projectionBindingExpression);
+                        var propertyMap = projectionIndex is IDictionary<IProperty, int>
+                            ? (IDictionary<IProperty, int>)projectionIndex
+                            : ((QueryableJsonProjectionInfo)projectionIndex).PropertyIndexMap;
+
                         _materializationContextBindings[parameterExpression] = propertyMap;
                         _entityTypeIdentifyingExpressionInfo[parameterExpression] =
                             // If single entity type is being selected in hierarchy then we use the value directly else we store the offset
@@ -535,6 +539,50 @@ protected override Expression VisitExtension(Expression extensionExpression)
                                 visitedShaperResultParameter,
                                 shaper.Type);
                         }
+                        else if (GetProjectionIndex(projectionBindingExpression) is QueryableJsonProjectionInfo queryableJsonEntityProjectionInfo)
+                        {
+                            if (_isTracking)
+                            {
+                                throw new InvalidOperationException(
+                                    RelationalStrings.JsonEntityOrCollectionProjectedAtRootLevelInTrackingQuery(nameof(EntityFrameworkQueryableExtensions.AsNoTracking)));
+                            }
+
+                            // json entity converted to query root and projected
+                            var entityParameter = Parameter(shaper.Type);
+                            _variables.Add(entityParameter);
+                            var entityMaterializationExpression = (BlockExpression)_parentVisitor.InjectEntityMaterializers(shaper);
+
+                            var mappedProperties = queryableJsonEntityProjectionInfo.PropertyIndexMap.Keys.ToList();
+                            var rewrittenEntityMaterializationExpression = new QueryableJsonEntityMaterializerRewriter(mappedProperties)
+                                .Rewrite(entityMaterializationExpression);
+
+                            var visitedEntityMaterializationExpression = Visit(rewrittenEntityMaterializationExpression);
+                            _expressions.Add(Assign(entityParameter, visitedEntityMaterializationExpression));
+
+                            foreach (var childProjectionInfo in queryableJsonEntityProjectionInfo.ChildrenProjectionInfo)
+                            {
+                                var (jsonReaderDataVariable, keyValuesParameter) = JsonShapingPreProcess(
+                                    childProjectionInfo.JsonProjectionInfo,
+                                    childProjectionInfo.Navigation.TargetEntityType,
+                                    childProjectionInfo.Navigation.IsCollection);
+
+                                var shaperResult = CreateJsonShapers(
+                                    childProjectionInfo.Navigation.TargetEntityType,
+                                    nullable: true,
+                                    jsonReaderDataVariable,
+                                    keyValuesParameter,
+                                    parentEntityExpression: entityParameter,
+                                    navigation: childProjectionInfo.Navigation);
+
+                                var visitedShaperResult = Visit(shaperResult);
+
+                                _includeExpressions.Add(visitedShaperResult);
+                            }
+
+                            accessor = CompensateForCollectionMaterialization(
+                                entityParameter,
+                                shaper.Type);
+                        }
                         else
                         {
                             var entityParameter = Parameter(shaper.Type);
@@ -2141,6 +2189,62 @@ ParameterExpression ExtractAndCacheNonConstantJsonArrayElementAccessValue(int in
             }
         }
 
+        private sealed class QueryableJsonEntityMaterializerRewriter : ExpressionVisitor
+        {
+            private readonly List<IProperty> _mappedProperties;
+
+            public QueryableJsonEntityMaterializerRewriter(List<IProperty> mappedProperties)
+            {
+                _mappedProperties = mappedProperties;
+            }
+
+            public BlockExpression Rewrite(BlockExpression jsonEntityShaperMaterializer)
+                => (BlockExpression)VisitBlock(jsonEntityShaperMaterializer);
+
+            protected override Expression VisitBinary(BinaryExpression binaryExpression)
+            {
+                // here we try to pattern match part of the shaper code that checks if key values are null
+                // if they are all non-null then we generate the entity
+                // problem for JSON entities is that some of the keys are synthesized and should be omitted
+                // if the key is one of the mapped ones, we leave the expression as is, otherwise replace with Constant(true)
+                // i.e. removing it
+                if (binaryExpression is
+                    {
+                        NodeType: ExpressionType.NotEqual,
+                        Left: MethodCallExpression
+                        {
+                            Method: { IsGenericMethod: true } method,
+                            Arguments: [_, _, ConstantExpression { Value: IProperty property }]
+                        },
+                        Right: ConstantExpression { Value: null }
+                    }
+                    && method.GetGenericMethodDefinition() == Infrastructure.ExpressionExtensions.ValueBufferTryReadValueMethod)
+                {
+                    return _mappedProperties.Contains(property)
+                        ? binaryExpression
+                        : Constant(true);
+                }
+
+                return base.VisitBinary(binaryExpression);
+            }
+
+            protected override Expression VisitMethodCall(MethodCallExpression methodCallExpression)
+            {
+                if (methodCallExpression is
+                    {
+                        Method: { IsGenericMethod: true } method,
+                        Arguments: [_, _, ConstantExpression { Value: IProperty property }]
+                    }
+                    && method.GetGenericMethodDefinition() == Infrastructure.ExpressionExtensions.ValueBufferTryReadValueMethod
+                    && !_mappedProperties.Contains(property))
+                {
+                    return Default(methodCallExpression.Type);
+                }
+
+                return base.VisitMethodCall(methodCallExpression);
+            }
+        }
+
         private static LambdaExpression GenerateFixup(
             Type entityType,
             Type relatedEntityType,