Correctly set source for chained group by operators

Author: hazzik

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

@@ -847,6 +847,8 @@ public void NestedGroupBy()
 				.GroupBy(x => new { x.Order.OrderId, x.Product.ProductId }) // this works fine
 				.GroupBy(x => x.Key.ProductId) // exception: "A recognition error occurred"
 				.ToList();
+
+			Assert.That(list, Has.Count.EqualTo(77));
 		}
 
 		private static void CheckGrouping<TKey, TElement>(IEnumerable<IGrouping<TKey, TElement>> groupedItems, Func<TElement, TKey> groupBy)

src/NHibernate/Linq/GroupBy/NonAggregatingGroupByRewriter.cs

@@ -2,7 +2,6 @@
 using System.Linq;
 using System.Linq.Expressions;
 using NHibernate.Linq.ResultOperators;
-using NHibernate.Util;
 using Remotion.Linq;
 using Remotion.Linq.Clauses;
 using Remotion.Linq.Clauses.ExpressionVisitors;
@@ -18,13 +17,12 @@ public static void ReWrite(QueryModel queryModel)
 			if (queryModel.ResultOperators.All(r => r is GroupResultOperator)
 			    && IsNonAggregatingGroupBy(queryModel))
 			{
-				var resultOperators = queryModel.ResultOperators
-					.ToArray(r => new NonAggregatingGroupBy((GroupResultOperator) r));
-				
-				queryModel.ResultOperators.Clear();
-				foreach (var resultOperator in resultOperators)
+				var source = queryModel.SelectClause.Selector;
+				for (var index = 0; index < queryModel.ResultOperators.Count; index++)
 				{
-					queryModel.ResultOperators.Add(resultOperator);
+					var r = (GroupResultOperator) queryModel.ResultOperators[index];
+					queryModel.ResultOperators[index] = new NonAggregatingGroupBy(r, source);
+					source = new QuerySourceReferenceExpression(r);
 				}
 
 				return;
@@ -65,7 +63,7 @@ private static void FlattenSubQuery(SubQueryExpression subQueryExpression, Query
 				throw new NotImplementedException();
 			}
 
-			queryModel.ResultOperators.Add(new NonAggregatingGroupBy((GroupResultOperator) subQueryModel.ResultOperators[0]));
+			queryModel.ResultOperators.Add(new NonAggregatingGroupBy((GroupResultOperator) subQueryModel.ResultOperators[0], queryModel.SelectClause.Selector));
 			queryModel.ResultOperators.Add(clientSideSelect);
 		}
 

src/NHibernate/Linq/ResultOperators/NonAggregatingGroupBy.cs

@@ -1,14 +1,25 @@
+using System;
+using System.Linq.Expressions;
 using Remotion.Linq.Clauses.ResultOperators;
 
 namespace NHibernate.Linq.ResultOperators
 {
 	public class NonAggregatingGroupBy : ClientSideTransformOperator
 	{
-		public NonAggregatingGroupBy(GroupResultOperator groupBy)
+		//Since v5.4
+		[Obsolete("Please use another constructor")]
+		public NonAggregatingGroupBy(GroupResultOperator groupBy) : this(groupBy, null)
+		{
+		}
+
+		public NonAggregatingGroupBy(GroupResultOperator groupBy, Expression source)
 		{
 			GroupBy = groupBy;
+			Source = source;
 		}
 
-		public GroupResultOperator GroupBy { get; private set; }
+		public GroupResultOperator GroupBy { get; }
+
+		public Expression Source { get; }
 	}
-}
+}

src/NHibernate/Linq/Visitors/ResultOperatorProcessors/ProcessNonAggregatingGroupBy.cs

@@ -11,7 +11,7 @@ public class ProcessNonAggregatingGroupBy : IResultOperatorProcessor<NonAggregat
 	{
 		public void Process(NonAggregatingGroupBy resultOperator, QueryModelVisitor queryModelVisitor, IntermediateHqlTree tree)
 		{
-			var selector = queryModelVisitor.Model.SelectClause.Selector;
+			var selector = resultOperator.Source;
 			var keySelector = resultOperator.GroupBy.KeySelector;
 			var elementSelector = resultOperator.GroupBy.ElementSelector;