Merge pull request #170 from DmitryNaumov/NH-3320

Fix NH-3320

Author: oskarb

src/NHibernate.Test/Linq/SelectionTests.cs

@@ -320,6 +320,14 @@ public void CanProjectWithCast()
 			// ReSharper restore RedundantCast
 		}
 
+		[Test]
+		public void CanSelectAfterOrderByAndTake()
+		{
+			// NH-3320
+			var names = db.Users.OrderBy(p => p.Name).Take(3).Select(p => p.Name).ToList();
+			Assert.AreEqual(3, names.Count);
+		}
+
 		[Test]
 		public void CanSelectManyWithCast()
 		{

src/NHibernate/Linq/GroupBy/PagingRewriter.cs

@@ -0,0 +1,68 @@
+using System.Linq;
+using NHibernate.Linq.Visitors;
+using Remotion.Linq;
+using Remotion.Linq.Clauses;
+using Remotion.Linq.Clauses.Expressions;
+using Remotion.Linq.Clauses.ResultOperators;
+
+namespace NHibernate.Linq.GroupBy
+{
+	public static class PagingRewriter
+	{
+		private static readonly System.Type[] PagingResultOperators = new[]
+																		  {
+																			  typeof (SkipResultOperator),
+																			  typeof (TakeResultOperator),
+																		  };
+
+		public static void ReWrite(QueryModel queryModel)
+		{
+			var subQueryExpression = queryModel.MainFromClause.FromExpression as SubQueryExpression;
+
+			if (subQueryExpression != null &&
+				subQueryExpression.QueryModel.ResultOperators.All(x => PagingResultOperators.Contains(x.GetType())))
+			{
+				FlattenSubQuery(subQueryExpression, queryModel);
+			}
+		}
+
+		private static void FlattenSubQuery(SubQueryExpression subQueryExpression, QueryModel queryModel)
+		{
+			// we can not flattern subquery if outer query has body clauses.
+			var subQueryModel = subQueryExpression.QueryModel;
+			var subQueryMainFromClause = subQueryModel.MainFromClause;
+			if (queryModel.BodyClauses.Count == 0)
+			{
+				foreach (var resultOperator in subQueryModel.ResultOperators)
+					queryModel.ResultOperators.Add(resultOperator);
+
+				foreach (var bodyClause in subQueryModel.BodyClauses)
+					queryModel.BodyClauses.Add(bodyClause);
+			}
+			else
+			{
+				var cro = new ContainsResultOperator(new QuerySourceReferenceExpression(subQueryMainFromClause));
+
+				var newSubQueryModel = subQueryModel.Clone();
+				newSubQueryModel.ResultOperators.Add(cro);
+				newSubQueryModel.ResultTypeOverride = typeof(bool);
+
+				var where = new WhereClause(new SubQueryExpression(newSubQueryModel));
+				queryModel.BodyClauses.Add(where);
+
+				if (!queryModel.BodyClauses.OfType<OrderByClause>().Any())
+				{
+					var orderByClauses = subQueryModel.BodyClauses.OfType<OrderByClause>();
+					foreach (var orderByClause in orderByClauses)
+						queryModel.BodyClauses.Add(orderByClause);
+				}
+			}
+
+			// Point all query source references to the outer from clause
+			queryModel.TransformExpressions(s => new SwapQuerySourceVisitor(queryModel.MainFromClause, subQueryMainFromClause).Swap(s));
+
+			// Replace the outer query source
+			queryModel.MainFromClause = subQueryMainFromClause;
+		}
+	}
+}

src/NHibernate/Linq/Visitors/QueryModelVisitor.cs

@@ -37,6 +37,9 @@ public static ExpressionToHqlTranslationResults GenerateHqlQuery(QueryModel quer
 			// Rewrite non-aggregating group-joins
 			NonAggregatingGroupJoinRewriter.ReWrite(queryModel);
 
+			// Rewrite paging
+			PagingRewriter.ReWrite(queryModel);
+
 			// Flatten pointless subqueries
 			QueryReferenceExpressionFlattener.ReWrite(queryModel);
 

src/NHibernate/NHibernate.csproj

@@ -288,9 +288,10 @@
     <Compile Include="Linq\Clauses\NhHavingClause.cs" />
     <Compile Include="Linq\Clauses\NhJoinClause.cs" />
     <Compile Include="Linq\ExpressionTransformers\RemoveCharToIntConversion.cs" />
-	<Compile Include="Linq\ExpressionTransformers\RemoveRedundantCast.cs" />
+    <Compile Include="Linq\ExpressionTransformers\RemoveRedundantCast.cs" />
     <Compile Include="Linq\Functions\DictionaryGenerator.cs" />
     <Compile Include="Linq\Functions\EqualsGenerator.cs" />
+    <Compile Include="Linq\GroupBy\PagingRewriter.cs" />
     <Compile Include="Linq\Visitors\NhPartialEvaluatingExpressionTreeVisitor.cs" />
     <Compile Include="Linq\QuerySourceNamer.cs" />
     <Compile Include="Linq\ReWriters\AddJoinsReWriter.cs" />