From 15cbebbbf9d8c0dc038c0e859714e64624a9713d Mon Sep 17 00:00:00 2001
From: Marco Belladelli <marcobladel@gmail.com>
Date: Thu, 30 May 2024 12:05:45 +0200
Subject: [PATCH] HHH-18121 Replace dynamic instantiations for non-simple count
 queries

---
 .../sqm/tree/select/SqmSelectStatement.java   | 37 ++++++++++++++-----
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/hibernate-core/src/main/java/org/hibernate/query/sqm/tree/select/SqmSelectStatement.java b/hibernate-core/src/main/java/org/hibernate/query/sqm/tree/select/SqmSelectStatement.java
index 55ec360b10..c92455af9d 100644
--- a/hibernate-core/src/main/java/org/hibernate/query/sqm/tree/select/SqmSelectStatement.java
+++ b/hibernate-core/src/main/java/org/hibernate/query/sqm/tree/select/SqmSelectStatement.java
@@ -483,16 +483,7 @@ public SqmSelectStatement<Long> createCountQuery() {
 			return (SqmSelectStatement<Long>) copy;
 		}
 		else {
-			final JpaSelection<?> selection = queryPart.getFirstQuerySpec().getSelection();
-			if ( selection.isCompoundSelection() ) {
-				char c = 'a';
-				for ( JpaSelection<?> item : selection.getSelectionItems() ) {
-					item.alias( Character.toString( ++c ) + '_' );
-				}
-			}
-			else {
-				selection.alias( "a_" );
-			}
+			aliasSelections( queryPart );
 			final SqmSubQuery<?> subquery = new SqmSubQuery<>( copy, queryPart, null, nodeBuilder() );
 			final SqmSelectStatement<Long> query = nodeBuilder().createQuery( Long.class );
 			query.from( subquery );
@@ -500,4 +491,30 @@ public SqmSelectStatement<Long> createCountQuery() {
 			return query;
 		}
 	}
+
+	private <S> void aliasSelections(SqmQueryPart<S> queryPart) {
+		if ( queryPart.isSimpleQueryPart() ) {
+			final SqmQuerySpec<S> querySpec = queryPart.getFirstQuerySpec();
+			final LinkedHashSet<JpaSelection<?>> newSelections = new LinkedHashSet<>();
+			aliasSelection( querySpec.getSelection(), newSelections );
+			//noinspection unchecked
+			querySpec.setSelection( (JpaSelection<S>) ( newSelections.size() == 1 ?
+					newSelections.iterator().next() :
+					nodeBuilder().tuple( newSelections.toArray( new JpaSelection<?>[0] ) ) ) );
+		}
+		else {
+			( (SqmQueryGroup<?>) queryPart ).getQueryParts().forEach( this::aliasSelections );
+		}
+	}
+
+	private void aliasSelection(JpaSelection<?> selection, LinkedHashSet<JpaSelection<?>> newSelections) {
+		if ( selection.isCompoundSelection() || selection instanceof SqmDynamicInstantiation<?> ) {
+			for ( JpaSelection<?> selectionItem : selection.getSelectionItems() ) {
+				aliasSelection( selectionItem, newSelections );
+			}
+		}
+		else {
+			newSelections.add( selection.alias( "c" + newSelections.size() ) );
+		}
+	}
 }