Merge pull request #432 from PleasantD/NH-3797

NH-3797 - Use group by keys as first class HqlCandidates for select clause rewriting

Author: hazzik

src/NHibernate.Test/Linq/ByMethod/GroupByTests.cs

@@ -528,6 +528,27 @@ into grp
 			Assert.That(result.Count, Is.EqualTo(77));
 		}
 
+		[Test(Description = "NH-3797")]
+		public void GroupByComputedValue()
+		{
+			var orderGroups = db.Orders.GroupBy(o => o.Customer.CustomerId == null ? 0 : 1).Select(g => new { Key = g.Key, Count = g.Count() }).ToList();
+			Assert.AreEqual(830, orderGroups.Sum(g => g.Count));
+		}
+
+		[Test(Description = "NH-3797")]
+		public void GroupByComputedValueInAnonymousType()
+		{
+			var orderGroups = db.Orders.GroupBy(o => new { Key = o.Customer.CustomerId == null ? 0 : 1 }).Select(g => new { Key = g.Key, Count = g.Count() }).ToList();
+			Assert.AreEqual(830, orderGroups.Sum(g => g.Count));
+		}
+
+		[Test(Description = "NH-3797")]
+		public void GroupByComputedValueInObjectArray()
+		{
+			var orderGroups = db.Orders.GroupBy(o => new[] { o.Customer.CustomerId == null ? 0 : 1, }).Select(g => new { Key = g.Key, Count = g.Count() }).ToList();
+			Assert.AreEqual(830, orderGroups.Sum(g => g.Count));
+		}
+
 		private static void CheckGrouping<TKey, TElement>(IEnumerable<IGrouping<TKey, TElement>> groupedItems, Func<TElement, TKey> groupBy)
 		{
 			var used = new HashSet<object>();

src/NHibernate/Linq/GroupBy/AggregatingGroupByRewriter.cs

@@ -1,9 +1,11 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Linq.Expressions;
 using NHibernate.Linq.Clauses;
 using NHibernate.Linq.ReWriters;
 using NHibernate.Linq.Visitors;
+using NHibernate.Util;
 using Remotion.Linq;
 using Remotion.Linq.Clauses.Expressions;
 using Remotion.Linq.Clauses.ResultOperators;
@@ -42,7 +44,7 @@ public static class AggregatingGroupByRewriter
 				typeof (CacheableResultOperator)
 			};
 
-		public static void ReWrite(QueryModel queryModel)
+		public static void ReWrite(QueryModel queryModel, IList<Expression> groupByKeys)
 		{
 			var subQueryExpression = queryModel.MainFromClause.FromExpression as SubQueryExpression;
 
@@ -57,6 +59,7 @@ public static void ReWrite(QueryModel queryModel)
 					var groupBy = operators[0] as GroupResultOperator;
 					if (groupBy != null)
 					{
+						groupBy.ExtractKeyExpressions(groupByKeys);
 						FlattenSubQuery(queryModel, subQueryExpression.QueryModel, groupBy);
 					}
 				}
@@ -91,7 +94,7 @@ private static void FlattenSubQuery(QueryModel queryModel, QueryModel subQueryMo
 				queryModel.BodyClauses.Add(bodyClause);
 
 			// Replace the outer select clause...
-			queryModel.SelectClause.TransformExpressions(s => 
+			queryModel.SelectClause.TransformExpressions(s =>
 				GroupBySelectClauseRewriter.ReWrite(s, groupBy, subQueryModel));
 
 			// Point all query source references to the outer from clause

src/NHibernate/Linq/GroupResultOperatorExtensions.cs

@@ -0,0 +1,23 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Linq.Expressions;
+using System.Text;
+using NHibernate.Util;
+using Remotion.Linq.Clauses.ResultOperators;
+
+namespace NHibernate.Linq
+{
+	internal static class GroupResultOperatorExtensions
+	{
+		public static void ExtractKeyExpressions(this GroupResultOperator groupResult, IList<Expression> groupByKeys)
+		{
+			if (groupResult.KeySelector is NewExpression)
+				(groupResult.KeySelector as NewExpression).Arguments.ForEach(groupByKeys.Add);
+			else if (groupResult.KeySelector is NewArrayExpression)
+				(groupResult.KeySelector as NewArrayExpression).Expressions.ForEach(groupByKeys.Add);
+			else
+				groupByKeys.Add(groupResult.KeySelector);
+		}
+	}
+}

src/NHibernate/Linq/Visitors/QueryModelVisitor.cs

@@ -1,4 +1,6 @@
 using System;
+using System.Collections;
+using System.Collections.Generic;
 using System.Linq.Expressions;
 using NHibernate.Hql.Ast;
 using NHibernate.Linq.Clauses;
@@ -32,7 +34,7 @@ public static ExpressionToHqlTranslationResults GenerateHqlQuery(QueryModel quer
 			NonAggregatingGroupByRewriter.ReWrite(queryModel);
 
 			// Rewrite aggregate group-by statements
-			AggregatingGroupByRewriter.ReWrite(queryModel);
+			AggregatingGroupByRewriter.ReWrite(queryModel, parameters.GroupByKeys);
 
 			// Rewrite aggregating group-joins
 			AggregatingGroupJoinRewriter.ReWrite(queryModel);
@@ -77,7 +79,10 @@ public static ExpressionToHqlTranslationResults GenerateHqlQuery(QueryModel quer
 			// Identify and name query sources
 			QuerySourceIdentifier.Visit(parameters.QuerySourceNamer, queryModel);
 
-			var visitor = new QueryModelVisitor(parameters, root, queryModel) { RewrittenOperatorResult = result };
+			var visitor = new QueryModelVisitor(parameters, root, queryModel)
+			{
+				RewrittenOperatorResult = result,
+			};
 			visitor.Visit();
 
 			return visitor._hqlTree.GetTranslation();
@@ -233,7 +238,7 @@ public override void VisitSelectClause(SelectClause selectClause, QueryModel que
 		{
 			CurrentEvaluationType = selectClause.GetOutputDataInfo();
 
-			var visitor = new SelectClauseVisitor(typeof(object[]), VisitorParameters);
+			var visitor = new SelectClauseVisitor(typeof(object[]), VisitorParameters, VisitorParameters.GroupByKeys);
 
 			visitor.Visit(selectClause.Selector);
 

src/NHibernate/Linq/Visitors/SelectClauseVisitor.cs

@@ -18,11 +18,12 @@ public class SelectClauseVisitor : ExpressionTreeVisitor
 		private List<HqlExpression> _hqlTreeNodes = new List<HqlExpression>();
 		private readonly HqlGeneratorExpressionTreeVisitor _hqlVisitor;
 
-		public SelectClauseVisitor(System.Type inputType, VisitorParameters parameters)
+		public SelectClauseVisitor(System.Type inputType, VisitorParameters parameters, IEnumerable<Expression> groupByKeys)
 		{
 			_inputParameter = Expression.Parameter(inputType, "input");
 			_parameters = parameters;
 			_hqlVisitor = new HqlGeneratorExpressionTreeVisitor(_parameters);
+			_hqlNodes = new HashSet<Expression>(groupByKeys);
 		}
 
 		public LambdaExpression ProjectionExpression { get; private set; }
@@ -43,7 +44,7 @@ public void Visit(Expression expression)
 			// Find the sub trees that can be expressed purely in HQL
 			var nominator = new SelectClauseHqlNominator(_parameters);
 			nominator.Visit(expression);
-			_hqlNodes = nominator.HqlCandidates;
+			_hqlNodes.UnionWith(nominator.HqlCandidates);
 
 			// Linq2SQL ignores calls to local methods. Linq2EF seems to not support
 			// calls to local methods at all. For NHibernate we support local methods,

src/NHibernate/Linq/Visitors/VisitorParameters.cs

@@ -16,6 +16,8 @@ public class VisitorParameters
 
 		public QuerySourceNamer QuerySourceNamer { get; set; }
 
+		public IList<Expression> GroupByKeys { get; private set; }
+
 		public VisitorParameters(
 			ISessionFactoryImplementor sessionFactory, 
 			IDictionary<ConstantExpression, NamedParameter> constantToParameterMap, 
@@ -26,6 +28,7 @@ public VisitorParameters(
 			ConstantToParameterMap = constantToParameterMap;
 			RequiredHqlParameters = requiredHqlParameters;
 			QuerySourceNamer = querySourceNamer;
+			GroupByKeys = new List<Expression>();
 		}
 	}
 }

src/NHibernate/NHibernate.csproj

@@ -299,6 +299,7 @@
     <Compile Include="Linq\Functions\EqualsGenerator.cs" />
     <Compile Include="Linq\GroupBy\KeySelectorVisitor.cs" />
     <Compile Include="Linq\GroupBy\PagingRewriter.cs" />
+    <Compile Include="Linq\GroupResultOperatorExtensions.cs" />
     <Compile Include="Linq\NestedSelects\NestedSelectDetector.cs" />
     <Compile Include="Linq\NestedSelects\Tuple.cs" />
     <Compile Include="Linq\NestedSelects\SelectClauseRewriter.cs" />