Undo not working code

Author: bahusoid

src/NHibernate/Linq/Visitors/JoinBuilder.cs

@@ -4,6 +4,7 @@
 using Remotion.Linq;
 using Remotion.Linq.Clauses;
 using Remotion.Linq.Clauses.Expressions;
+using Remotion.Linq.Parsing;
 
 namespace NHibernate.Linq.Visitors
 {
@@ -75,5 +76,23 @@ private void AddJoin(QueryModel queryModel, NhJoinClause joinClause)
 		{
 			queryModel.BodyClauses.Add(joinClause);
 		}
+
+		private class QuerySourceExtractor : RelinqExpressionVisitor
+		{
+			private IQuerySource _querySource;
+
+			public static IQuerySource GetQuerySource(Expression expression)
+			{
+				var sourceExtractor = new QuerySourceExtractor();
+				sourceExtractor.Visit(expression);
+				return sourceExtractor._querySource;
+			}
+
+			protected override Expression VisitQuerySourceReference(QuerySourceReferenceExpression expression)
+			{
+				_querySource = expression.ReferencedQuerySource;
+				return base.VisitQuerySourceReference(expression);
+			}
+		}
 	}
 }

src/NHibernate/Linq/Visitors/QuerySourceExtractor.cs

@@ -6,8 +6,6 @@
 namespace NHibernate.Linq.Visitors
 {
 
-	// Since 5.5
-	[Obsolete("This class has no more usages and will be removed in a future version")]
 	public class QuerySourceLocator : NhQueryModelVisitorBase
 	{
 		readonly System.Type _type;

src/NHibernate/Linq/Visitors/QuerySourceLocator.cs

@@ -4,7 +4,6 @@
 using NHibernate.Hql.Ast.ANTLR;
 using NHibernate.Persister.Entity;
 using NHibernate.Type;
-using Remotion.Linq.Clauses.StreamedData;
 using Remotion.Linq.EagerFetching;
 
 namespace NHibernate.Linq.Visitors.ResultOperatorProcessors
@@ -13,17 +12,18 @@ public class ProcessFetch
 	{
 		public void Process(FetchRequestBase resultOperator, QueryModelVisitor queryModelVisitor, IntermediateHqlTree tree)
 		{
-			var selector = ((StreamedSequenceInfo) queryModelVisitor.PreviousEvaluationType).ItemExpression;
-			var querySource = QuerySourceExtractor.GetQuerySource(selector);
-			var sourceAlias = queryModelVisitor.VisitorParameters.QuerySourceNamer.GetName(querySource);
-			Process(resultOperator, queryModelVisitor, tree, GetFromNodeByAlias(tree, sourceAlias), sourceAlias);
+			var querySource = QuerySourceLocator.FindQuerySource(
+				queryModelVisitor.Model,
+				resultOperator.RelationMember.DeclaringType);
+			var name = queryModelVisitor.VisitorParameters.QuerySourceNamer.GetName(querySource);
+			Process(resultOperator, queryModelVisitor, tree, GetFromNodeByAlias(tree, name), name);
 		}
 
-		private static HqlTreeNode GetFromNodeByAlias(IntermediateHqlTree tree, string sourceAlias)
+		private static HqlTreeNode GetFromNodeByAlias(IntermediateHqlTree tree, string alias)
 		{
 			return tree.Root.NodesPreOrder
 				.First(x => x.AstNode.Type == HqlSqlWalker.FROM).Children
-				.First(x => GetNodeAlias(x) == sourceAlias);
+				.First(x => GetNodeAlias(x) == alias);
 		}
 
 		private static string GetNodeAlias(HqlTreeNode fromNode)