Cache only required parts of NhLinqExpression in QueryExpressionPlan

Author: hazzik

src/NHibernate/Engine/Query/QueryExpressionPlan.cs

@@ -1,7 +1,6 @@
 using System;
 using System.Collections.Generic;
 using NHibernate.Hql;
-using NHibernate.Linq;
 
 namespace NHibernate.Engine.Query
 {

src/NHibernate/Engine/Query/QueryPlanCache.cs

@@ -60,10 +60,11 @@ public IQueryExpressionPlan GetHQLQueryPlan(IQueryExpression queryExpression, bo
 				{
 					log.Debug("unable to locate HQL query plan in cache; generating ({0})", queryExpression.Key);
 				}
+
 				plan = new QueryExpressionPlan(queryExpression, shallow, enabledFilters, factory);
 				// 6.0 TODO: add "CanCachePlan { get; }" to IQueryExpression interface
 				if (!(queryExpression is ICacheableQueryExpression linqExpression) || linqExpression.CanCachePlan)
-					planCache.Put(key, plan);
+					planCache.Put(key, PreparePlanToCache(plan));
 				else
 					log.Debug("Query plan not cacheable");
 			}
@@ -79,23 +80,30 @@ public IQueryExpressionPlan GetHQLQueryPlan(IQueryExpression queryExpression, bo
 			return plan;
 		}
 
+		private QueryExpressionPlan PreparePlanToCache(QueryExpressionPlan plan)
+		{
+			if (plan.QueryExpression is NhLinqExpression planExpression)
+			{
+				return plan.Copy(new NhLinqExpressionCache(planExpression));
+			}
+
+			return plan;
+		}
+
 		private static QueryExpressionPlan CopyIfRequired(QueryExpressionPlan plan, IQueryExpression queryExpression)
 		{
-			var planExpression = plan.QueryExpression as NhLinqExpression;
-			var expression = queryExpression as NhLinqExpression;
-			if (planExpression != null && expression != null)
+			if (plan.QueryExpression is NhLinqExpressionCache cache && queryExpression is NhLinqExpression expression)
 			{
 				//NH-3413
 				//Here we have to use original expression.
 				//In most cases NH do not translate expression in second time, but 
 				// for cases when we have list parameters in query, like @p1.Contains(...),
 				// it does, and then it uses parameters from first try. 
-				//TODO: cache only required parts of QueryExpression
 
 				//NH-3436
 				// We have to return new instance plan with it's own query expression
-				// because other treads can override queryexpression of current plan during execution of query if we will use cached instance of plan 
-				expression.CopyExpressionTranslation(planExpression);
+				// because other treads can override query expression of current plan during execution of query if we will use cached instance of plan 
+				expression.CopyExpressionTranslation(cache);
 				plan = plan.Copy(expression);
 			}
 
@@ -118,7 +126,7 @@ public IQueryExpressionPlan GetFilterQueryPlan(IQueryExpression queryExpression,
 				plan = new FilterQueryPlan(queryExpression, collectionRole, shallow, enabledFilters, factory);
 				// 6.0 TODO: add "CanCachePlan { get; }" to IQueryExpression interface
 				if (!(queryExpression is ICacheableQueryExpression linqExpression) || linqExpression.CanCachePlan)
-					planCache.Put(key, plan);
+					planCache.Put(key, PreparePlanToCache(plan));
 				else
 					log.Debug("Query plan not cacheable");
 			}

src/NHibernate/Hql/Ast/ANTLR/ASTQueryTranslatorFactory.cs

@@ -33,7 +33,7 @@ static IQueryTranslator[] CreateQueryTranslators(
 
 			var translators = polymorphicParsers
 				.ToArray(hql => queryExpression is NhLinqExpression linqExpression
-							? new QueryTranslatorImpl(queryIdentifier, hql, filters, factory, linqExpression.NamedParameters)
+							? new QueryTranslatorImpl(queryIdentifier, hql, filters, factory, linqExpression.GetNamedParameterTypes())
 							: new QueryTranslatorImpl(queryIdentifier, hql, filters, factory));
 
 			foreach (var translator in translators)

src/NHibernate/Hql/Ast/ANTLR/HqlSqlWalker.cs

@@ -55,7 +55,6 @@ public partial class HqlSqlWalker
 		private readonly LiteralProcessor _literalProcessor;
 
 		private readonly IDictionary<string, string> _tokenReplacements;
-		private readonly IDictionary<string, NamedParameter> _namedParameters;
 		private readonly IDictionary<IParameterSpecification, IType> _guessedParameterTypes = new Dictionary<IParameterSpecification, IType>();
 
 		private JoinType _impliedJoinType;
@@ -65,32 +64,20 @@ public partial class HqlSqlWalker
 		private IASTFactory _nodeFactory;
 		private readonly List<AssignmentSpecification> assignmentSpecifications = new List<AssignmentSpecification>();
 		private int numberOfParametersInSetClause;
-		private Stack<int> clauseStack=new Stack<int>();
+		private Stack<int> clauseStack = new Stack<int>();
 
 		public HqlSqlWalker(
 			QueryTranslatorImpl qti,
 			ISessionFactoryImplementor sfi,
 			ITreeNodeStream input,
 			IDictionary<string, string> tokenReplacements,
 			string collectionRole)
-			: this(qti, sfi, input, tokenReplacements, null, collectionRole)
-		{
-		}
-
-		internal HqlSqlWalker(
-			QueryTranslatorImpl qti,
-			ISessionFactoryImplementor sfi,
-			ITreeNodeStream input,
-			IDictionary<string, string> tokenReplacements,
-			IDictionary<string, NamedParameter> namedParameters,
-			string collectionRole)
 			: this(input)
 		{
 			_sessionFactoryHelper = new SessionFactoryHelperExtensions(sfi);
 			_qti = qti;
 			_literalProcessor = new LiteralProcessor(this);
 			_tokenReplacements = tokenReplacements;
-			_namedParameters = namedParameters;
 			_collectionFilterRole = collectionRole;
 		}
 
@@ -313,8 +300,7 @@ void PostProcessInsert(IASTNode insert)
 
 				IASTNode idSelectExprNode = null;
 
-				var seqGenerator = generator as SequenceGenerator;
-				if (seqGenerator != null)
+				if (generator is SequenceGenerator seqGenerator)
 				{
 					string seqName = seqGenerator.GeneratorKey();
 					string nextval = SessionFactoryHelper.Factory.Dialect.GetSelectSequenceNextValString(seqName);
@@ -1083,17 +1069,17 @@ IASTNode GenerateNamedParameter(IASTNode delimiterNode, IASTNode nameNode)
 			);
 
 			parameter.HqlParameterSpecification = paramSpec;
-			if (_namedParameters != null && _namedParameters.TryGetValue(name, out var namedParameter))
+			if (_qti.TryGetNamedParameterType(name, out var type, out var isGuessedType))
 			{
 				// Add the parameter type information so that we are able to calculate functions return types
 				// when the parameter is used as an argument.
-				if (namedParameter.IsGuessedType)
+				if (isGuessedType)
 				{
-					_guessedParameterTypes[paramSpec] = namedParameter.Type;
-					parameter.GuessedType = namedParameter.Type;
+					_guessedParameterTypes[paramSpec] = type;
+					parameter.GuessedType = type;
 				}
 				else
-					parameter.ExpectedType = namedParameter.Type;
+					parameter.ExpectedType = type;
 			}
 
 			_parameters.Add(paramSpec);

src/NHibernate/Hql/Ast/ANTLR/QueryTranslatorImpl.cs

@@ -29,7 +29,7 @@ public partial class QueryTranslatorImpl : IFilterTranslator
 		private readonly string _queryIdentifier;
 		private readonly IASTNode _stageOneAst;
 		private readonly ISessionFactoryImplementor _factory;
-		private readonly IDictionary<string, NamedParameter> _namedParameters;
+		private readonly IDictionary<string, Tuple<IType, bool>> _namedParameters;
 
 		private bool _shallowQuery;
 		private bool _compiled;
@@ -69,7 +69,7 @@ internal QueryTranslatorImpl(
 			IASTNode parsedQuery,
 			IDictionary<string, IFilter> enabledFilters,
 			ISessionFactoryImplementor factory,
-			IDictionary<string, NamedParameter> namedParameters)
+			IDictionary<string, Tuple<IType, bool>> namedParameters)
 		{
 			_queryIdentifier = queryIdentifier;
 			_stageOneAst = parsedQuery;
@@ -454,7 +454,7 @@ private static IStatementExecutor BuildAppropriateStatementExecutor(IStatement s
 
 		private HqlSqlTranslator Analyze(string collectionRole)
 		{
-			var translator = new HqlSqlTranslator(_stageOneAst, this, _factory, _tokenReplacements, _namedParameters, collectionRole);
+			var translator = new HqlSqlTranslator(_stageOneAst, this, _factory, _tokenReplacements, collectionRole);
 
 			translator.Translate();
 
@@ -468,6 +468,20 @@ private void ErrorIfDML()
 				throw new QueryExecutionRequestException("Not supported for DML operations", _queryIdentifier);
 			}
 		}
+
+		public bool TryGetNamedParameterType(string name, out IType type, out bool isGuessedType)
+		{
+			if (_namedParameters == null || !_namedParameters.TryGetValue(name, out var p))
+			{
+				type = null;
+				isGuessedType = false;
+				return false;
+			}
+
+			type = p.Item1;
+			isGuessedType = p.Item2;
+			return true;
+		}
 	}
 
 	public class HqlParseEngine
@@ -568,23 +582,20 @@ internal class HqlSqlTranslator
 		private readonly QueryTranslatorImpl _qti;
 		private readonly ISessionFactoryImplementor _sfi;
 		private readonly IDictionary<string, string> _tokenReplacements;
-		private readonly IDictionary<string, NamedParameter> _namedParameters;
 		private readonly string _collectionRole;
 		private IStatement _resultAst;
 
-		public HqlSqlTranslator(
+		internal HqlSqlTranslator(
 			IASTNode ast,
 			QueryTranslatorImpl qti,
 			ISessionFactoryImplementor sfi,
 			IDictionary<string, string> tokenReplacements,
-			IDictionary<string, NamedParameter> namedParameters,
 			string collectionRole)
 		{
 			_inputAst = ast;
 			_qti = qti;
 			_sfi = sfi;
 			_tokenReplacements = tokenReplacements;
-			_namedParameters = namedParameters;
 			_collectionRole = collectionRole;
 		}
 
@@ -604,7 +615,7 @@ public IStatement Translate()
 
 				var nodes = new BufferedTreeNodeStream(_inputAst);
 
-				var hqlSqlWalker = new HqlSqlWalker(_qti, _sfi, nodes, _tokenReplacements, _namedParameters, _collectionRole);
+				var hqlSqlWalker = new HqlSqlWalker(_qti, _sfi, nodes, _tokenReplacements, _collectionRole);
 				hqlSqlWalker.TreeAdaptor = new HqlSqlWalkerTreeAdaptor(hqlSqlWalker);
 
 				try

src/NHibernate/Linq/NhLinqExpression.cs

@@ -113,12 +113,18 @@ public IASTNode Translate(ISessionFactoryImplementor sessionFactory, bool filter
 			return DuplicateTree(ExpressionToHqlTranslationResults.Statement.AstNode);
 		}
 
-		internal void CopyExpressionTranslation(NhLinqExpression other)
+		internal void CopyExpressionTranslation(NhLinqExpressionCache cache)
 		{
-			ExpressionToHqlTranslationResults = other.ExpressionToHqlTranslationResults;
-			ParameterDescriptors = other.ParameterDescriptors;
+			ExpressionToHqlTranslationResults = cache.ExpressionToHqlTranslationResults;
+			ParameterDescriptors = cache.ParameterDescriptors;
 			// Type could have been overridden by translation.
-			Type = other.Type;
+			Type = cache.Type;
+		}
+
+		internal IDictionary<string, Tuple<IType, bool>> GetNamedParameterTypes()
+		{
+			return _constantToParameterMap.Values.Distinct()
+				.ToDictionary(p => p.Name, p => System.Tuple.Create(p.Type, p.IsGuessedType));
 		}
 
 		private static IASTNode DuplicateTree(IASTNode ast)

src/NHibernate/Linq/NhLinqExpressionCache.cs

@@ -0,0 +1,29 @@
+using System;
+using System.Collections.Generic;
+using NHibernate.Engine;
+using NHibernate.Engine.Query;
+using NHibernate.Hql.Ast.ANTLR.Tree;
+
+namespace NHibernate.Linq
+{
+	internal class NhLinqExpressionCache : IQueryExpression
+	{
+		internal NhLinqExpressionCache(NhLinqExpression expression)
+		{
+			ExpressionToHqlTranslationResults = expression.ExpressionToHqlTranslationResults ?? throw new ArgumentException("NhLinqExpression is not translated");
+			Key = expression.Key;
+			Type = expression.Type;
+			ParameterDescriptors = expression.ParameterDescriptors;
+		}
+
+		public ExpressionToHqlTranslationResults ExpressionToHqlTranslationResults { get; }
+		public string Key { get; }
+		public System.Type Type { get; }
+		public IList<NamedParameterDescriptor> ParameterDescriptors { get; }
+
+		public IASTNode Translate(ISessionFactoryImplementor sessionFactory, bool filter)
+		{
+			return ExpressionToHqlTranslationResults.Statement.AstNode;
+		}
+	}
+}