Query: Add support for GroupBy patterns beyond aggregate

src/EFCore.InMemory/Query/Internal/EntityProjectionExpression.cs

@@ -148,6 +148,27 @@ public virtual void AddNavigationBinding(INavigation navigation, EntityShaperExp
                 : null;
         }
 
+        /// <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 virtual EntityProjectionExpression Clone()
+        {
+            var readExpressionMap = new Dictionary<IProperty, MethodCallExpression>(_readExpressionMap);
+            var entityProjectionExpression = new EntityProjectionExpression(EntityType, readExpressionMap);
+            foreach (var kvp in _navigationExpressionsCache)
+            {
+                entityProjectionExpression._navigationExpressionsCache[kvp.Key] = new EntityShaperExpression(
+                    kvp.Value.EntityType,
+                    ((EntityProjectionExpression)kvp.Value.ValueBufferExpression).Clone(),
+                    kvp.Value.IsNullable);
+            }
+
+            return entityProjectionExpression;
+        }
+
         /// <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

src/EFCore.InMemory/Query/Internal/InMemoryQueryExpression.Helper.cs

@@ -173,5 +173,64 @@ protected override Expression VisitExtension(Expression extensionExpression)
                     : base.VisitExtension(extensionExpression);
             }
         }
+
+        private sealed class QueryExpressionReplacingExpressionVisitor : ExpressionVisitor
+        {
+            private readonly Expression _oldQuery;
+            private readonly Expression _newQuery;
+
+            public QueryExpressionReplacingExpressionVisitor(Expression oldQuery, Expression newQuery)
+            {
+                _oldQuery = oldQuery;
+                _newQuery = newQuery;
+            }
+
+            [return: NotNullIfNotNull("expression")]
+            public override Expression? Visit(Expression? expression)
+            {
+                return expression is ProjectionBindingExpression projectionBindingExpression
+                    && ReferenceEquals(projectionBindingExpression.QueryExpression, _oldQuery)
+                    ? projectionBindingExpression.ProjectionMember != null
+                        ? new ProjectionBindingExpression(
+                            _newQuery, projectionBindingExpression.ProjectionMember!, projectionBindingExpression.Type)
+                        : new ProjectionBindingExpression(
+                            _newQuery, projectionBindingExpression.Index!.Value, projectionBindingExpression.Type)
+                    : base.Visit(expression);
+            }
+        }
+
+        private sealed class CloningExpressionVisitor : ExpressionVisitor
+        {
+            [return: NotNullIfNotNull("expression")]
+            public override Expression? Visit(Expression? expression)
+            {
+                if (expression is InMemoryQueryExpression inMemoryQueryExpression)
+                {
+                    var clonedInMemoryQueryExpression = new InMemoryQueryExpression(
+                        inMemoryQueryExpression.ServerQueryExpression, inMemoryQueryExpression._valueBufferParameter)
+                    {
+                        _groupingParameter = inMemoryQueryExpression._groupingParameter,
+                        _singleResultMethodInfo = inMemoryQueryExpression._singleResultMethodInfo,
+                        _scalarServerQuery = inMemoryQueryExpression._scalarServerQuery
+                    };
+
+                    clonedInMemoryQueryExpression._clientProjections.AddRange(inMemoryQueryExpression._clientProjections.Select(e => Visit(e)));
+                    clonedInMemoryQueryExpression._projectionMappingExpressions.AddRange(inMemoryQueryExpression._projectionMappingExpressions);
+                    foreach (var item in inMemoryQueryExpression._projectionMapping)
+                    {
+                        clonedInMemoryQueryExpression._projectionMapping[item.Key] = Visit(item.Value);
+                    }
+
+                    return clonedInMemoryQueryExpression;
+                }
+
+                if (expression is EntityProjectionExpression entityProjectionExpression)
+                {
+                    return entityProjectionExpression.Clone();
+                }
+
+                return base.Visit(expression);
+            }
+        }
     }
 }

src/EFCore.InMemory/Query/Internal/InMemoryQueryExpression.cs

@@ -41,10 +41,20 @@ private static readonly ConstructorInfo _resultEnumerableConstructor
         private MethodInfo? _singleResultMethodInfo;
         private bool _scalarServerQuery;
 
+        private CloningExpressionVisitor? _cloningExpressionVisitor;
+
         private Dictionary<ProjectionMember, Expression> _projectionMapping = new();
         private readonly List<Expression> _clientProjections = new();
         private readonly List<Expression> _projectionMappingExpressions = new();
 
+        private InMemoryQueryExpression(
+            Expression serverQueryExpression,
+            ParameterExpression valueBufferParameter)
+        {
+            ServerQueryExpression = serverQueryExpression;
+            _valueBufferParameter = valueBufferParameter;
+        }
+
         /// <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
@@ -333,15 +343,17 @@ public virtual void ApplyProjection()
                 ServerQueryExpression,
                 selectorLambda);
 
+            _groupingParameter = null;
+
             if (_singleResultMethodInfo != null)
             {
                 ServerQueryExpression = Call(
                     _singleResultMethodInfo.MakeGenericMethod(CurrentParameter.Type),
                     ServerQueryExpression);
 
-                _singleResultMethodInfo = null;
-
                 ConvertToEnumerable();
+
+                _singleResultMethodInfo = null;
             }
         }
 
@@ -540,7 +552,7 @@ public virtual void ApplyDistinct()
         ///     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 virtual InMemoryGroupByShaperExpression ApplyGrouping(
+        public virtual GroupByShaperExpression ApplyGrouping(
             Expression groupingKey,
             Expression shaperExpression,
             bool defaultElementSelector)
@@ -583,11 +595,15 @@ public virtual InMemoryGroupByShaperExpression ApplyGrouping(
                 keySelector,
                 selector);
 
-            return new InMemoryGroupByShaperExpression(
+            var clonedInMemoryQueryExpression = Clone();
+            clonedInMemoryQueryExpression.UpdateServerQueryExpression(_groupingParameter);
+            clonedInMemoryQueryExpression._groupingParameter = null;
+
+            return new GroupByShaperExpression(
                 groupingKey,
-                shaperExpression,
-                _groupingParameter,
-                _valueBufferParameter);
+                new ShapedQueryExpression(
+                    clonedInMemoryQueryExpression,
+                    new QueryExpressionReplacingExpressionVisitor(this, clonedInMemoryQueryExpression).Visit(shaperExpression)));
         }
 
         /// <summary>
@@ -711,6 +727,22 @@ public virtual EntityShaperExpression AddNavigationToWeakEntityType(
             return entityShaper;
         }
 
+        /// <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 virtual ShapedQueryExpression Clone(Expression shaperExpression)
+        {
+            var clonedInMemoryQueryExpression = Clone();
+
+            return new ShapedQueryExpression(
+                clonedInMemoryQueryExpression,
+                new QueryExpressionReplacingExpressionVisitor(this, clonedInMemoryQueryExpression).Visit(shaperExpression));
+
+        }
+
         /// <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
@@ -811,6 +843,16 @@ void IPrintableExpression.Print(ExpressionPrinter expressionPrinter)
             }
         }
 
+        private InMemoryQueryExpression Clone()
+        {
+            if (_cloningExpressionVisitor == null)
+            {
+                _cloningExpressionVisitor = new();
+            }
+
+            return (InMemoryQueryExpression)_cloningExpressionVisitor.Visit(this);
+        }
+
         private Expression GetGroupingKey(Expression key, List<Expression> groupingExpressions, Expression groupingKeyAccessExpression)
         {
             switch (key)
@@ -1061,7 +1103,7 @@ static Expression MakeNullable(Expression expression, bool nullable)
 
         private void ConvertToEnumerable()
         {
-            if (ServerQueryExpression.Type.TryGetSequenceType() == null)
+            if (_scalarServerQuery || _singleResultMethodInfo != null)
             {
                 if (ServerQueryExpression.Type != typeof(ValueBuffer))
                 {

src/EFCore.InMemory/Query/Internal/InMemoryQueryableMethodTranslatingExpressionVisitor.cs

@@ -70,6 +70,24 @@ protected InMemoryQueryableMethodTranslatingExpressionVisitor(
         protected override QueryableMethodTranslatingExpressionVisitor CreateSubqueryVisitor()
             => new InMemoryQueryableMethodTranslatingExpressionVisitor(this);
 
+        /// <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>
+        protected override Expression VisitExtension(Expression extensionExpression)
+        {
+            if (extensionExpression is GroupByShaperExpression groupByShaperExpression)
+            {
+                var shapedQueryExpression = groupByShaperExpression.GroupingEnumerable;
+                return ((InMemoryQueryExpression)shapedQueryExpression.QueryExpression)
+                    .Clone(shapedQueryExpression.ShaperExpression);
+            }
+
+            return base.VisitExtension(extensionExpression);
+        }
+
         /// <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
@@ -1513,7 +1531,8 @@ private static Expression AccessField(
                         inMemoryQueryExpression.CurrentParameter)
                     : TranslateLambdaExpression(source, selector, preserveType: true);
 
-            if (selector == null)
+            if (selector == null
+                || selector.Body is EntityProjectionExpression)
             {
                 return null;
             }