Use early exit in QueryTransformers and introspectors.

Avoid iterating tokens when things had been discovered already and reduce visibility of StringQuery to previous value.

See: #3309

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/EqlCountQueryTransformer.java

@@ -15,7 +15,9 @@
  */
 package org.springframework.data.jpa.repository.query;
 
-import static org.springframework.data.jpa.repository.query.QueryTokens.*;
+import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_CLOSE_PAREN;
+import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_COMMA;
+import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_COUNT_FUNC;
 
 import org.springframework.data.jpa.repository.query.QueryRenderer.QueryRendererBuilder;
 import org.springframework.data.jpa.repository.query.QueryTransformers.CountSelectionTokenStream;
@@ -27,6 +29,7 @@ import org.springframework.lang.Nullable;
  *
  * @author Greg Turnquist
  * @author Mark Paluch
+ * @author Christoph Strobl
  * @since 3.4
  */
 @SuppressWarnings("ConstantValue")
@@ -36,6 +39,7 @@ class EqlCountQueryTransformer extends EqlQueryRenderer {
 	private final @Nullable String primaryFromAlias;
 
 	EqlCountQueryTransformer(@Nullable String countProjection, @Nullable String primaryFromAlias) {
+
 		this.countProjection = countProjection;
 		this.primaryFromAlias = primaryFromAlias;
 	}
@@ -64,47 +68,45 @@ class EqlCountQueryTransformer extends EqlQueryRenderer {
 	@Override
 	public QueryTokenStream visitSelect_clause(EqlParser.Select_clauseContext ctx) {
 
+		boolean usesDistinct = ctx.DISTINCT() != null;
+
 		QueryRendererBuilder builder = QueryRenderer.builder();
 
 		builder.append(QueryTokens.expression(ctx.SELECT()));
 		builder.append(TOKEN_COUNT_FUNC);
 
-		if (countProjection != null) {
-			builder.append(QueryTokens.token(countProjection));
-		}
-
 		QueryRendererBuilder nested = QueryRenderer.builder();
-
-		if (ctx.DISTINCT() != null) {
-			nested.append(QueryTokens.expression(ctx.DISTINCT()));
-		}
-
 		if (countProjection == null) {
-
-			if (ctx.DISTINCT() != null) {
-
-				QueryTokenStream selectionListbuilder = QueryTokenStream.concat(ctx.select_item(), this::visit,
-						TOKEN_COMMA);
-
-				CountSelectionTokenStream countSelection = QueryTransformers
-						.filterCountSelection(selectionListbuilder);
-
-				if (countSelection.requiresPrimaryAlias()) {
-					// constructor
-					nested.append(QueryTokens.token(primaryFromAlias));
-				} else {
-					// keep all the select items to distinct against
-					nested.append(countSelection);
-				}
+			if (usesDistinct) {
+				nested.append(QueryTokens.expression(ctx.DISTINCT()));
+				nested.append(getDistinctCountSelection(QueryTokenStream.concat(ctx.select_item(), this::visit, TOKEN_COMMA)));
 			} else {
 				nested.append(QueryTokens.token(primaryFromAlias));
 			}
+		} else {
+			builder.append(QueryTokens.token(countProjection));
+			if (usesDistinct) {
+				nested.append(QueryTokens.expression(ctx.DISTINCT()));
+			}
 		}
 
 		builder.appendInline(nested);
 		builder.append(TOKEN_CLOSE_PAREN);
-
 		return builder;
 	}
 
+	private QueryRendererBuilder getDistinctCountSelection(QueryTokenStream selectionListbuilder) {
+
+		QueryRendererBuilder nested = new QueryRendererBuilder();
+		CountSelectionTokenStream countSelection = QueryTransformers.filterCountSelection(selectionListbuilder);
+
+		if (countSelection.requiresPrimaryAlias()) {
+			// constructor
+			nested.append(QueryTokens.token(primaryFromAlias));
+		} else {
+			// keep all the select items to distinct against
+			nested.append(countSelection);
+		}
+		return nested;
+	}
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/EqlQueryIntrospector.java

@@ -15,18 +15,20 @@
  */
 package org.springframework.data.jpa.repository.query;
 
-import static org.springframework.data.jpa.repository.query.QueryTokens.*;
+import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_COMMA;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
+import org.springframework.data.jpa.repository.query.EqlParser.Range_variable_declarationContext;
 import org.springframework.lang.Nullable;
 
 /**
  * {@link ParsedQueryIntrospector} for EQL queries.
  *
  * @author Mark Paluch
+ * @author Christoph Strobl
  */
 @SuppressWarnings("UnreachableCode")
 class EqlQueryIntrospector extends EqlBaseVisitor<Void> implements ParsedQueryIntrospector {
@@ -56,20 +58,8 @@ class EqlQueryIntrospector extends EqlBaseVisitor<Void> implements ParsedQueryIn
 	@Override
 	public Void visitSelect_clause(EqlParser.Select_clauseContext ctx) {
 
-		List<EqlParser.Select_itemContext> selections = ctx.select_item();
-		List<QueryToken> selectItemTokens = new ArrayList<>(selections.size() * 2);
-
-		for (EqlParser.Select_itemContext selection : selections) {
-
-			if (!selectItemTokens.isEmpty()) {
-				selectItemTokens.add(TOKEN_COMMA);
-			}
-
-			selectItemTokens.add(QueryTokens.token(QueryRenderer.from(renderer.visitSelect_item(selection)).render()));
-		}
-
 		if (!projectionProcessed) {
-			projection = selectItemTokens;
+			projection = captureSelectItems(ctx.select_item(), renderer);
 			projectionProcessed = true;
 		}
 
@@ -80,8 +70,7 @@ class EqlQueryIntrospector extends EqlBaseVisitor<Void> implements ParsedQueryIn
 	public Void visitRange_variable_declaration(EqlParser.Range_variable_declarationContext ctx) {
 
 		if (primaryFromAlias == null) {
-			primaryFromAlias = ctx.identification_variable() != null ? ctx.identification_variable().getText()
-					: ctx.entity_name().getText();
+			primaryFromAlias = capturePrimaryAlias(ctx);
 		}
 
 		return super.visitRange_variable_declaration(ctx);
@@ -91,8 +80,26 @@ class EqlQueryIntrospector extends EqlBaseVisitor<Void> implements ParsedQueryIn
 	public Void visitConstructor_expression(EqlParser.Constructor_expressionContext ctx) {
 
 		hasConstructorExpression = true;
-
 		return super.visitConstructor_expression(ctx);
 	}
 
+	private static String capturePrimaryAlias(Range_variable_declarationContext ctx) {
+		return ctx.identification_variable() != null ? ctx.identification_variable().getText()
+				: ctx.entity_name().getText();
+	}
+
+	private static List<QueryToken> captureSelectItems(List<EqlParser.Select_itemContext> selections,
+			EqlQueryRenderer itemRenderer) {
+
+		List<QueryToken> selectItemTokens = new ArrayList<>(selections.size() * 2);
+		for (EqlParser.Select_itemContext selection : selections) {
+
+			if (!selectItemTokens.isEmpty()) {
+				selectItemTokens.add(TOKEN_COMMA);
+			}
+
+			selectItemTokens.add(QueryTokens.token(QueryRenderer.from(itemRenderer.visitSelect_item(selection)).render()));
+		}
+		return selectItemTokens;
+	}
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/HqlCountQueryTransformer.java

@@ -15,8 +15,14 @@
  */
 package org.springframework.data.jpa.repository.query;
 
-import static org.springframework.data.jpa.repository.query.QueryTokens.*;
+import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_AS;
+import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_CLOSE_PAREN;
+import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_COUNT_FUNC;
+import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_DOUBLE_UNDERSCORE;
+import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_OPEN_PAREN;
+import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_SELECT_COUNT;
 
+import org.springframework.data.jpa.repository.query.HqlParser.SelectClauseContext;
 import org.springframework.data.jpa.repository.query.QueryRenderer.QueryRendererBuilder;
 import org.springframework.data.jpa.repository.query.QueryTransformers.CountSelectionTokenStream;
 import org.springframework.lang.Nullable;
@@ -175,53 +181,30 @@ class HqlCountQueryTransformer extends HqlQueryRenderer {
 		QueryRendererBuilder builder = QueryRenderer.builder();
 		builder.append(QueryTokens.expression(ctx.SELECT()));
 
-		QueryTokenStream selectionListbuilder = visit(ctx.selectionList());
-
-		if (!isSubquery(ctx)) {
-
-			builder.append(TOKEN_COUNT_FUNC);
-
-			if (countProjection != null) {
-				builder.append(QueryTokens.token(countProjection));
-			}
-
-			QueryRendererBuilder nested = QueryRenderer.builder();
-
-			if (ctx.DISTINCT() != null) {
-				nested.append(QueryTokens.expression(ctx.DISTINCT()));
-			}
-
-			if (countProjection == null) {
-
-				if (ctx.DISTINCT() != null) {
-
-					CountSelectionTokenStream countSelection = QueryTransformers
-							.filterCountSelection(selectionListbuilder);
-
-					if (countSelection.requiresPrimaryAlias()) {
-						// constructor
-						nested.append(QueryTokens.token(primaryFromAlias));
-					} else {
-						// keep all the select items to distinct against
-						nested.append(countSelection);
-					}
-				} else {
-					nested.append(QueryTokens.token(primaryFromAlias));
-				}
-			}
-
-			builder.appendInline(nested);
-			builder.append(TOKEN_CLOSE_PAREN);
-
-		} else {
-
-			if (ctx.DISTINCT() != null) {
-				builder.append(QueryTokens.expression(ctx.DISTINCT()));
-			}
-
-			builder.append(selectionListbuilder);
+		if (isSubquery(ctx)) {
+			return visitSubQuerySelectClause(ctx, builder);
 		}
 
+		builder.append(TOKEN_COUNT_FUNC);
+		boolean usesDistinct = ctx.DISTINCT() != null;
+		QueryRendererBuilder nested = QueryRenderer.builder();
+		if (countProjection == null) {
+			if (usesDistinct) {
+
+				nested.append(QueryTokens.expression(ctx.DISTINCT()));
+				nested.append(getDistinctCountSelection(visit(ctx.selectionList())));
+			} else {
+				nested.append(QueryTokens.token(primaryFromAlias));
+			}
+		} else {
+			builder.append(QueryTokens.token(countProjection));
+			if (usesDistinct) {
+				nested.append(QueryTokens.expression(ctx.DISTINCT()));
+			}
+		}
+
+		builder.appendInline(nested);
+		builder.append(TOKEN_CLOSE_PAREN);
 		return builder;
 	}
 
@@ -245,4 +228,28 @@ class HqlCountQueryTransformer extends HqlQueryRenderer {
 		return builder;
 	}
 
+	private QueryRendererBuilder visitSubQuerySelectClause(SelectClauseContext ctx, QueryRendererBuilder builder) {
+		if (ctx.DISTINCT() != null) {
+			builder.append(QueryTokens.expression(ctx.DISTINCT()));
+		}
+
+		builder.append(visit(ctx.selectionList()));
+		return builder;
+	}
+
+	private QueryRendererBuilder getDistinctCountSelection(QueryTokenStream selectionListbuilder) {
+
+		QueryRendererBuilder nested = new QueryRendererBuilder();
+		CountSelectionTokenStream countSelection = QueryTransformers.filterCountSelection(selectionListbuilder);
+
+		if (countSelection.requiresPrimaryAlias()) {
+			// constructor
+			nested.append(QueryTokens.token(primaryFromAlias));
+		} else {
+			// keep all the select items to distinct against
+			nested.append(countSelection);
+		}
+		return nested;
+	}
+
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/HqlQueryIntrospector.java

@@ -15,12 +15,13 @@
  */
 package org.springframework.data.jpa.repository.query;
 
-import static org.springframework.data.jpa.repository.query.QueryTokens.*;
+import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_COMMA;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
+import org.springframework.data.jpa.repository.query.HqlParser.VariableContext;
 import org.springframework.lang.Nullable;
 
 /**
@@ -56,20 +57,8 @@ class HqlQueryIntrospector extends HqlBaseVisitor<Void> implements ParsedQueryIn
 	@Override
 	public Void visitSelectClause(HqlParser.SelectClauseContext ctx) {
 
-		List<HqlParser.SelectionContext> selections = ctx.selectionList().selection();
-		List<QueryToken> selectItemTokens = new ArrayList<>(selections.size() * 2);
-
-		for (HqlParser.SelectionContext selection : selections) {
-
-			if (!selectItemTokens.isEmpty()) {
-				selectItemTokens.add(TOKEN_COMMA);
-			}
-
-			selectItemTokens.add(QueryTokens.token(QueryRenderer.from(renderer.visitSelection(selection)).render()));
-		}
-
 		if (!projectionProcessed) {
-			projection = selectItemTokens;
+			projection = captureSelectItems(ctx.selectionList().selection(), renderer);
 			projectionProcessed = true;
 		}
 
@@ -80,12 +69,10 @@ class HqlQueryIntrospector extends HqlBaseVisitor<Void> implements ParsedQueryIn
 	public Void visitFromRoot(HqlParser.FromRootContext ctx) {
 
 		if (primaryFromAlias == null && ctx.variable() != null && !HqlQueryRenderer.isSubquery(ctx)) {
-
-			primaryFromAlias = (ctx.variable().reservedWord() != null ? ctx.variable().reservedWord()
-					: ctx.variable().identifier().reservedWord()).getText();
+			primaryFromAlias = capturePrimaryAlias(ctx.variable());
 		}
 
-		return null;
+		return super.visitFromRoot(ctx);
 	}
 
 	@Override
@@ -96,4 +83,22 @@ class HqlQueryIntrospector extends HqlBaseVisitor<Void> implements ParsedQueryIn
 		return super.visitInstantiation(ctx);
 	}
 
+	private static String capturePrimaryAlias(VariableContext ctx) {
+		return ((ctx).reservedWord() != null ? ctx.reservedWord() : ctx.identifier().reservedWord()).getText();
+	}
+
+	private static List<QueryToken> captureSelectItems(List<HqlParser.SelectionContext> selections,
+			HqlQueryRenderer itemRenderer) {
+
+		List<QueryToken> selectItemTokens = new ArrayList<>(selections.size() * 2);
+		for (HqlParser.SelectionContext selection : selections) {
+
+			if (!selectItemTokens.isEmpty()) {
+				selectItemTokens.add(TOKEN_COMMA);
+			}
+
+			selectItemTokens.add(QueryTokens.token(QueryRenderer.from(itemRenderer.visitSelection(selection)).render()));
+		}
+		return selectItemTokens;
+	}
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpqlCountQueryTransformer.java

@@ -27,6 +27,7 @@ import org.springframework.lang.Nullable;
  *
  * @author Greg Turnquist
  * @author Mark Paluch
+ * @author Christoph Strobl
  * @since 3.1
  */
 @SuppressWarnings("ConstantValue")
@@ -64,41 +65,26 @@ class JpqlCountQueryTransformer extends JpqlQueryRenderer {
 	@Override
 	public QueryRendererBuilder visitSelect_clause(JpqlParser.Select_clauseContext ctx) {
 
+		boolean usesDistinct = ctx.DISTINCT() != null;
+
 		QueryRendererBuilder builder = QueryRenderer.builder();
 
 		builder.append(QueryTokens.expression(ctx.SELECT()));
 		builder.append(TOKEN_COUNT_FUNC);
 
-		if (countProjection != null) {
-			builder.append(QueryTokens.token(countProjection));
-		}
-
 		QueryRendererBuilder nested = QueryRenderer.builder();
-
-		if (ctx.DISTINCT() != null) {
-			nested.append(QueryTokens.expression(ctx.DISTINCT()));
-		}
-
 		if (countProjection == null) {
-
-			if (ctx.DISTINCT() != null) {
-
-				QueryTokenStream selectionListbuilder = QueryTokenStream.concat(ctx.select_item(), this::visit,
-						TOKEN_COMMA);
-
-				CountSelectionTokenStream countSelection = QueryTransformers
-						.filterCountSelection(selectionListbuilder);
-
-				if (countSelection.requiresPrimaryAlias()) {
-					// constructor
-					nested.append(QueryTokens.token(primaryFromAlias));
-				} else {
-					// keep all the select items to distinct against
-					nested.append(countSelection);
-				}
+			if (usesDistinct) {
+				nested.append(QueryTokens.expression(ctx.DISTINCT()));
+				nested.append(getDistinctCountSelection(QueryTokenStream.concat(ctx.select_item(), this::visit, TOKEN_COMMA)));
 			} else {
 				nested.append(QueryTokens.token(primaryFromAlias));
 			}
+		} else {
+			builder.append(QueryTokens.token(countProjection));
+			if (usesDistinct) {
+				nested.append(QueryTokens.expression(ctx.DISTINCT()));
+			}
 		}
 
 		builder.appendInline(nested);
@@ -107,4 +93,19 @@ class JpqlCountQueryTransformer extends JpqlQueryRenderer {
 		return builder;
 	}
 
+	private QueryRendererBuilder getDistinctCountSelection(QueryTokenStream selectionListbuilder) {
+
+		QueryRendererBuilder nested = new QueryRendererBuilder();
+		CountSelectionTokenStream countSelection = QueryTransformers.filterCountSelection(selectionListbuilder);
+
+		if (countSelection.requiresPrimaryAlias()) {
+			// constructor
+			nested.append(QueryTokens.token(primaryFromAlias));
+		} else {
+			// keep all the select items to distinct against
+			nested.append(countSelection);
+		}
+		return nested;
+	}
+
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpqlQueryIntrospector.java

@@ -15,7 +15,7 @@
  */
 package org.springframework.data.jpa.repository.query;
 
-import static org.springframework.data.jpa.repository.query.QueryTokens.*;
+import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_COMMA;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -27,6 +27,7 @@ import org.springframework.lang.Nullable;
  * {@link ParsedQueryIntrospector} for JPQL queries.
  *
  * @author Mark Paluch
+ * @author Christoph Strobl
  */
 @SuppressWarnings({ "UnreachableCode", "ConstantValue" })
 class JpqlQueryIntrospector extends JpqlBaseVisitor<Void> implements ParsedQueryIntrospector {
@@ -55,30 +56,17 @@ class JpqlQueryIntrospector extends JpqlBaseVisitor<Void> implements ParsedQuery
 	public Void visitRange_variable_declaration(JpqlParser.Range_variable_declarationContext ctx) {
 
 		if (primaryFromAlias == null) {
-			primaryFromAlias = ctx.identification_variable() == null ? ctx.entity_name().getText()
-					: ctx.identification_variable().getText();
+			primaryFromAlias = capturePrimaryAlias(ctx);
 		}
 
-		return null;
+		return super.visitRange_variable_declaration(ctx);
 	}
 
 	@Override
 	public Void visitSelect_clause(JpqlParser.Select_clauseContext ctx) {
 
-		List<JpqlParser.Select_itemContext> selections = ctx.select_item();
-		List<QueryToken> selectItemTokens = new ArrayList<>(selections.size() * 2);
-
-		for (JpqlParser.Select_itemContext selection : selections) {
-
-			if (!selectItemTokens.isEmpty()) {
-				selectItemTokens.add(TOKEN_COMMA);
-			}
-
-			selectItemTokens.add(QueryTokens.token(QueryRenderer.from(renderer.visitSelect_item(selection)).render()));
-		}
-
 		if (!projectionProcessed) {
-			projection = selectItemTokens;
+			projection = captureSelectItems(ctx.select_item(), renderer);
 			projectionProcessed = true;
 		}
 
@@ -93,4 +81,24 @@ class JpqlQueryIntrospector extends JpqlBaseVisitor<Void> implements ParsedQuery
 		return super.visitConstructor_expression(ctx);
 	}
 
+	private static String capturePrimaryAlias(JpqlParser.Range_variable_declarationContext ctx) {
+		return ctx.identification_variable() != null ? ctx.identification_variable().getText()
+				: ctx.entity_name().getText();
+	}
+
+	private static List<QueryToken> captureSelectItems(List<JpqlParser.Select_itemContext> selections,
+			JpqlQueryRenderer itemRenderer) {
+
+		List<QueryToken> selectItemTokens = new ArrayList<>(selections.size() * 2);
+		for (JpqlParser.Select_itemContext selection : selections) {
+
+			if (!selectItemTokens.isEmpty()) {
+				selectItemTokens.add(TOKEN_COMMA);
+			}
+
+			selectItemTokens.add(QueryTokens.token(QueryRenderer.from(itemRenderer.visitSelect_item(selection)).render()));
+		}
+		return selectItemTokens;
+	}
+
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/StringQuery.java

@@ -56,7 +56,7 @@ import org.springframework.util.StringUtils;
  * @author Greg Turnquist
  * @author Yuriy Tsarkov
  */
-public class StringQuery implements DeclaredQuery {
+class StringQuery implements DeclaredQuery {
 
 	private final String query;
 	private final List<ParameterBinding> bindings;