Query: Fix bug in IncludeCompiler where we would resolve an incorrect…

… QSRE for grouped results.

src/EFCore.Specification.Tests/Query/OwnedQueryTestBase.cs

@@ -52,6 +52,41 @@ public virtual void Query_for_leaf_type_loads_all_owned_navs()
             }
         }
 
+        [Fact]
+        public virtual void Query_when_group_by()
+        {
+            using (var context = CreateContext())
+            {
+                var people = context.Set<OwnedPerson>().GroupBy(op => op.Id).ToList();
+
+                Assert.Equal(4, people.Count);
+                Assert.True(people.SelectMany(p => p).All(p => p.PersonAddress != null));
+                Assert.True(people.SelectMany(p => p).OfType<Branch>().All(b => b.BranchAddress != null));
+                Assert.True(people.SelectMany(p => p).OfType<LeafA>().All(a => a.LeafAAddress != null));
+                Assert.True(people.SelectMany(p => p).OfType<LeafB>().All(b => b.LeafBAddress != null));
+            }
+        }
+
+        [Fact]
+        public virtual void Query_when_subquery()
+        {
+            using (var context = CreateContext())
+            {
+                var people 
+                    = context.Set<OwnedPerson>()
+                        .Distinct()
+                        .Take(5)
+                        .Select(op => new { op })
+                        .ToList();
+
+                Assert.Equal(4, people.Count);
+                Assert.True(people.All(p => p.op.PersonAddress != null));
+                Assert.True(people.Select(p => p.op).OfType<Branch>().All(b => b.BranchAddress != null));
+                Assert.True(people.Select(p => p.op).OfType<LeafA>().All(a => a.LeafAAddress != null));
+                Assert.True(people.Select(p => p.op).OfType<LeafB>().All(b => b.LeafBAddress != null));
+            }
+        }
+
         protected abstract DbContext CreateContext();
     }
 }

src/EFCore/Query/ExpressionVisitors/Internal/QuerySourceTracingExpressionVisitor.cs

@@ -58,23 +58,17 @@ protected override Expression VisitQuerySourceReference(QuerySourceReferenceExpr
                 }
                 else
                 {
-                    var fromClauseBase = expression.ReferencedQuerySource as FromClauseBase;
-
-                    if (fromClauseBase != null)
+                    if (expression.ReferencedQuerySource is FromClauseBase fromClauseBase)
                     {
                         Visit(fromClauseBase.FromExpression);
                     }
-
-                    var joinClause = expression.ReferencedQuerySource as JoinClause;
-
-                    if (joinClause != null)
+
+                    if (expression.ReferencedQuerySource is JoinClause joinClause)
                     {
                         Visit(joinClause.InnerSequence);
                     }
 
-                    var groupJoinClause = expression.ReferencedQuerySource as GroupJoinClause;
-
-                    if (groupJoinClause != null)
+                    if (expression.ReferencedQuerySource is GroupJoinClause groupJoinClause)
                     {
                         if (groupJoinClause.JoinClause.Equals(_targetQuerySource))
                         {

src/EFCore/Query/Internal/IncludeCompiler.cs

@@ -74,7 +74,7 @@ public virtual void CompileIncludes(
                 return;
             }
 
-            foreach (var includeLoadTree in CreateIncludeLoadTrees(queryModel.SelectClause.Selector))
+            foreach (var includeLoadTree in CreateIncludeLoadTrees(queryModel))
             {
                 includeLoadTree.Compile(
                     _queryCompilationContext,
@@ -111,8 +111,7 @@ public virtual void LogIgnoredIncludes()
             }
         }
 
-        private IEnumerable<IncludeLoadTree> CreateIncludeLoadTrees(
-            Expression targetExpression)
+        private IEnumerable<IncludeLoadTree> CreateIncludeLoadTrees(QueryModel queryModel)
         {
             var querySourceTracingExpressionVisitor
                 = _querySourceTracingExpressionVisitorFactory.Create();
@@ -126,7 +125,7 @@ var querySourceTracingExpressionVisitor
                 var querySourceReferenceExpression
                     = querySourceTracingExpressionVisitor
                         .FindResultQuerySourceReferenceExpression(
-                            targetExpression,
+                            queryModel.GetOutputExpression(),
                             includeResultOperator.QuerySource);
 
                 if (querySourceReferenceExpression == null

src/EFCore/Query/Internal/QueryModelExtensions.cs

@@ -1,13 +1,17 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Linq.Expressions;
 using JetBrains.Annotations;
+using Microsoft.EntityFrameworkCore.Internal;
 using Microsoft.EntityFrameworkCore.Query.ExpressionVisitors;
 using Remotion.Linq;
 using Remotion.Linq.Clauses;
 using Remotion.Linq.Clauses.Expressions;
+using Remotion.Linq.Clauses.ResultOperators;
 
 namespace Microsoft.EntityFrameworkCore.Query.Internal
 {
@@ -16,7 +20,42 @@ namespace Microsoft.EntityFrameworkCore.Query.Internal
     ///     directly from your code. This API may change or be removed in future releases.
     /// </summary>
     public static class QueryModelExtensions
-    {
+    {        
+        /// <summary>
+        ///     This API supports the Entity Framework Core infrastructure and is not intended to be used
+        ///     directly from your code. This API may change or be removed in future releases.
+        /// </summary>
+        public static Expression GetOutputExpression([NotNull] this QueryModel queryModel)
+        {
+            var outputExpression = queryModel.SelectClause.Selector;
+
+            var groupResultOperator
+                = queryModel.ResultOperators.OfType<GroupResultOperator>().LastOrDefault();
+
+            if (groupResultOperator != null)
+            {
+                outputExpression = groupResultOperator.ElementSelector;
+            }
+            else if (queryModel.SelectClause.Selector.Type.IsGrouping())
+            {
+                var subqueryExpression
+                    = (queryModel.SelectClause.Selector
+                        .TryGetReferencedQuerySource() as MainFromClause)?.FromExpression as SubQueryExpression;
+
+                var nestedGroupResultOperator
+                    = subqueryExpression?.QueryModel?.ResultOperators
+                        ?.OfType<GroupResultOperator>()
+                        .LastOrDefault();
+
+                if (nestedGroupResultOperator != null)
+                {
+                    outputExpression = nestedGroupResultOperator.ElementSelector;
+                }
+            }
+
+            return outputExpression;
+        }
+
         /// <summary>
         ///     This API supports the Entity Framework Core infrastructure and is not intended to be used
         ///     directly from your code. This API may change or be removed in future releases.

test/EFCore.SqlServer.FunctionalTests/Query/OwnedQuerySqlServerTest.cs

@@ -38,6 +38,18 @@ public override void Query_for_leaf_type_loads_all_owned_navs()
             AssertSql("");
         }
 
+        [Fact(Skip = "#8907")]
+        public override void Query_when_group_by()
+        {
+            base.Query_when_group_by();
+        }
+
+        [Fact(Skip = "#8907")]
+        public override void Query_when_subquery()
+        {
+            base.Query_when_subquery();
+        }
+
         protected override DbContext CreateContext() => _fixture.CreateContext();
 
         private void AssertSql(params string[] expected)

test/EFCore.Sqlite.FunctionalTests/Query/OwnedQuerySqliteTest.cs

@@ -32,6 +32,18 @@ public override void Query_for_leaf_type_loads_all_owned_navs()
             base.Query_for_leaf_type_loads_all_owned_navs();
         }
 
+        [Fact(Skip = "#8907")]
+        public override void Query_when_group_by()
+        {
+            base.Query_when_group_by();
+        }
+
+        [Fact(Skip = "#8907")]
+        public override void Query_when_subquery()
+        {
+            base.Query_when_subquery();
+        }
+
         protected override DbContext CreateContext() => _fixture.CreateContext();
     }
 }