HHH-18979 Add Restriction.apply and JpaCriteriaQuery.getRoot()

documentation/src/main/asciidoc/introduction/Interacting.adoc

@@ -919,7 +919,7 @@ It's even possible to transform a HQL query string to a criteria query, and modi
 ----
 HibernateCriteriaBuilder builder = sessionFactory.getCriteriaBuilder();
 var query = builder.createQuery("from Book where year(publicationDate) > 2000", Book.class);
-var root = (Root<Book>) query.getRootList().get(0);
+var root = query.getRoot(0, Book.class);
 query.where(builder.like(root.get(Book_.title), builder.literal("Hibernate%")));
 query.orderBy(builder.asc(root.get(Book_.title)), builder.desc(root.get(Book_.isbn)));
 List<Book> matchingBooks = session.createSelectionQuery(query).getResultList();

hibernate-core/src/main/java/org/hibernate/query/criteria/CriteriaDefinition.java

@@ -53,7 +53,7 @@
  *     List<Book> books
  *             = new CriteriaDefinition<>(sessionFactory, Book.class,
  *                     "from Book left join fetch authors where type = BOOK") {{
- *                 var book = (JpaRoot<Book>) getSelection();
+ *                 var book = getRoot(0, Book.class);
  *                 where(getRestriction(), like(book.get(Book_.title), "%Hibernate%"));
  *                 orderBy(desc(book.get(Book_.publicationDate)), asc(book.get(Book_.isbn)));
  *             }}
@@ -148,8 +148,10 @@ public CriteriaDefinition(CriteriaDefinition<R> template) {
 	 */
 	public CriteriaDefinition(CriteriaDefinition<?> template, Class<R> resultType) {
 		super( template.getCriteriaBuilder() );
-		query = ((SqmSelectStatement<?>) template.query)
-				.createCopy( SqmCopyContext.simpleContext(), resultType );
+		if ( !(template.query instanceof SqmSelectStatement<?> selectStatement) ) {
+			throw new IllegalArgumentException( "Not a SqmSelectStatement" );
+		}
+		query = selectStatement.createCopy( SqmCopyContext.simpleContext(), resultType );
 	}
 
 	public CriteriaDefinition(SessionFactory factory, Class<R> resultType) {
@@ -220,6 +222,11 @@ public JpaCriteriaQuery<R> restrict(Predicate predicate) {
 		return existing == null ? where( predicate ) : where( existing, predicate );
 	}
 
+	@Override
+	public HibernateCriteriaBuilder getCriteriaBuilder() {
+		return query.getCriteriaBuilder();
+	}
+
 	@Override
 	public JpaCriteriaQuery<R> select(Selection<? extends R> selection) {
 		return query.select(selection);
@@ -401,10 +408,20 @@ public FetchClauseType getFetchClauseType() {
 	}
 
 	@Override
-	public List<Root<?>> getRootList() {
+	public List<? extends JpaRoot<?>> getRootList() {
 		return query.getRootList();
 	}
 
+	@Override
+	public <E> JpaRoot<? extends E> getRoot(int position, Class<E> type) {
+		return query.getRoot( position, type );
+	}
+
+	@Override
+	public <E> JpaRoot<? extends E> getRoot(String alias, Class<E> type) {
+		return query.getRoot( alias, type );
+	}
+
 	@Override
 	public Collection<? extends JpaCteCriteria<?>> getCteCriterias() {
 		return query.getCteCriterias();

hibernate-core/src/main/java/org/hibernate/query/criteria/JpaCriteriaQuery.java

@@ -61,13 +61,33 @@ public interface JpaCriteriaQuery<T> extends CriteriaQuery<T>, JpaQueryableCrite
 	/**
 	 * Return the {@linkplain #getRoots() roots} as a list.
 	 */
-	List<Root<?>> getRootList();
+	List<? extends JpaRoot<?>> getRootList();
 
-	@Override
-	@SuppressWarnings("unchecked")
-	default List<Order> getOrderList() {
-		return (List) getQueryPart().getSortSpecifications();
-	}
+	/**
+	 * Get a {@linkplain Root query root} element at the given position
+	 * with the given type.
+	 *
+	 * @param position the position of this root element
+	 * @param type the type of the root entity
+	 *
+	 * @throws IllegalArgumentException if the root entity at the given
+	 *         position is not of the given type, or if there are not
+	 *         enough root entities in the query
+	 */
+	<E> JpaRoot<? extends E> getRoot(int position, Class<E> type);
+
+	/**
+	 * Get a {@linkplain Root query root} element with the given alias
+	 * and the given type.
+	 *
+	 * @param alias the identification variable of the root element
+	 * @param type the type of the root entity
+	 *
+	 * @throws IllegalArgumentException if the root entity with the
+	 *         given alias is not of the given type, or if there is
+	 *         no root entities with the given alias
+	 */
+	<E> JpaRoot<? extends E> getRoot(String alias, Class<E> type);
 
 	/**
 	 * {@inheritDoc}
@@ -132,4 +152,6 @@ default List<Order> getOrderList() {
 
 	@Override
 	<U> JpaSubQuery<U> subquery(EntityType<U> type);
+
+	HibernateCriteriaBuilder getCriteriaBuilder();
 }

hibernate-core/src/main/java/org/hibernate/query/criteria/JpaQueryStructure.java

@@ -66,9 +66,9 @@ public interface JpaQueryStructure<T> extends JpaQueryPart<T> {
 
 	List<? extends JpaExpression<?>> getGroupingExpressions();
 
-	JpaQueryStructure<T> setGroupingExpressions(List<? extends JpaExpression<?>> grouping);
+	JpaQueryStructure<T> setGroupingExpressions(List<? extends Expression<?>> grouping);
 
-	JpaQueryStructure<T> setGroupingExpressions(JpaExpression<?>... grouping);
+	JpaQueryStructure<T> setGroupingExpressions(Expression<?>... grouping);
 
 	JpaPredicate getGroupRestriction();
 

hibernate-core/src/main/java/org/hibernate/query/restriction/Restriction.java

@@ -5,13 +5,15 @@
 package org.hibernate.query.restriction;
 
 import jakarta.persistence.criteria.CriteriaBuilder;
+import jakarta.persistence.criteria.CriteriaQuery;
 import jakarta.persistence.criteria.Predicate;
 import jakarta.persistence.criteria.Root;
 import jakarta.persistence.metamodel.SingularAttribute;
 import org.hibernate.Incubating;
 import org.hibernate.Internal;
 import org.hibernate.query.Order;
 import org.hibernate.query.SelectionQuery;
+import org.hibernate.query.criteria.JpaCriteriaQuery;
 import org.hibernate.query.range.Range;
 
 import java.util.List;
@@ -77,6 +79,17 @@ default Restriction<X> and(Restriction<X> restriction) {
 	@Internal
 	Predicate toPredicate(Root<? extends X> root, CriteriaBuilder builder);
 
+	/**
+	 * Apply this restriction to the given root entity of the given
+	 * {@linkplain CriteriaQuery criteria query}.
+	 */
+	default void apply(CriteriaQuery<?> query, Root<? extends X> root) {
+		if ( !(query instanceof JpaCriteriaQuery<?> criteriaQuery) ) {
+			throw new IllegalArgumentException( "Not a JpaCriteriaQuery" );
+		}
+		query.where( query.getRestriction(), toPredicate( root, criteriaQuery.getCriteriaBuilder() ) );
+	}
+
 	/**
 	 * Restrict the allowed values of the given attribute to the given
 	 * {@linkplain Range range}.

hibernate-core/src/main/java/org/hibernate/query/sqm/NodeBuilder.java

@@ -908,19 +908,19 @@ <T> SqmJsonValueExpression<T> jsonValue(
 	SqmSelectStatement<Tuple> createTupleQuery();
 
 	@Override
-	<Y> JpaCompoundSelection<Y> construct(Class<Y> resultClass, Selection<?>[] selections);
+	<Y> JpaCompoundSelection<Y> construct(Class<Y> resultClass, Selection<?>... selections);
 
 	@Override
 	<Y> JpaCompoundSelection<Y> construct(Class<Y> resultClass, List<? extends JpaSelection<?>> arguments);
 
 	@Override
-	JpaCompoundSelection<Tuple> tuple(Selection<?>[] selections);
+	JpaCompoundSelection<Tuple> tuple(Selection<?>... selections);
 
 	@Override
 	JpaCompoundSelection<Tuple> tuple(List<Selection<?>> selections);
 
 	@Override
-	JpaCompoundSelection<Object[]> array(Selection<?>[] selections);
+	JpaCompoundSelection<Object[]> array(Selection<?>... selections);
 
 	@Override
 	JpaCompoundSelection<Object[]> array(List<Selection<?>> selections);

hibernate-core/src/main/java/org/hibernate/query/sqm/internal/AbstractSqmSelectionQuery.java

@@ -4,7 +4,6 @@
  */
 package org.hibernate.query.sqm.internal;
 
-import jakarta.persistence.criteria.Root;
 import org.hibernate.HibernateException;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.graph.spi.AppliedGraph;
@@ -150,9 +149,7 @@ public SelectionQuery<R> setOrder(Order<? super R> order) {
 	@Override
 	public SelectionQuery<R> addRestriction(Restriction<? super R> restriction) {
 		final SqmSelectStatement<R> selectStatement = getSqmSelectStatement().copy( noParamCopyContext() );
-		final Root<? extends R> root = (Root<? extends R>) selectStatement.getRootList().get( 0 );
-		selectStatement.where( selectStatement.getRestriction(),
-				restriction.toPredicate( root, selectStatement.nodeBuilder() ) );
+		restriction.apply( selectStatement, selectStatement.<R>getRoot( 0, getExpectedResultType() ) );
 		// TODO: when the QueryInterpretationCache can handle caching criteria queries,
 		//       simply cache the new SQM as if it were a criteria query, and remove this:
 		getQueryOptions().setQueryPlanCachingEnabled( false );

hibernate-core/src/main/java/org/hibernate/query/sqm/internal/KeyBasedPagination.java

@@ -111,7 +111,7 @@ private static <R> JpaCompoundSelection<KeyedResult<R>> keyedResultConstructor(
 		return builder.construct( resultClass, asList( selected, builder.construct(List.class, newItems ) ) );
 	}
 
-	@SuppressWarnings({"rawtypes", "unchecked"})
+	@SuppressWarnings("rawtypes")
 	private static <C extends Comparable<? super C>> SqmPredicate keyPredicate(
 			Expression<? extends C> key, C keyValue, SortDirection direction,
 			List<SqmPath<?>> previousKeys, List<Comparable<?>> keyValues,

hibernate-core/src/main/java/org/hibernate/query/sqm/internal/SqmCriteriaNodeBuilder.java

@@ -40,7 +40,6 @@
 import org.hibernate.internal.CoreMessageLogger;
 import org.hibernate.internal.SessionFactoryRegistry;
 import org.hibernate.internal.util.StringHelper;
-import org.hibernate.internal.util.collections.CollectionHelper;
 import org.hibernate.jpa.spi.JpaCompliance;
 import org.hibernate.metamodel.model.domain.DomainType;
 import org.hibernate.metamodel.model.domain.JpaMetamodel;
@@ -167,6 +166,7 @@
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 import static java.util.Arrays.asList;
+import static org.hibernate.internal.util.collections.CollectionHelper.determineProperSizing;
 import static org.hibernate.query.internal.QueryHelper.highestPrecedenceType;
 import static org.hibernate.query.sqm.TrimSpec.fromCriteriaTrimSpec;
 
@@ -872,7 +872,7 @@ public JpaSearchOrder desc(JpaCteCriteriaAttribute x, boolean nullsFirst) {
 	}
 
 	@Override
-	public JpaCompoundSelection<Tuple> tuple(Selection<?>[] selections) {
+	public JpaCompoundSelection<Tuple> tuple(Selection<?>... selections) {
 		return tuple( Arrays.asList( selections ) );
 	}
 
@@ -916,62 +916,69 @@ public <R> SqmTuple<R> tuple(SqmExpressible<R> tupleType, List<? extends SqmExpr
 	}
 
 	@Override
-	public JpaCompoundSelection<Object[]> array(Selection<?>[] selections) {
-		return array( Arrays.asList( selections ) );
+	public JpaCompoundSelection<Object[]> array(Selection<?>... selections) {
+		return array( Object[].class,
+				Arrays.stream( selections ).map( selection -> (SqmSelectableNode<?>) selection ).toList() );
 	}
 
 	@Override
 	public JpaCompoundSelection<Object[]> array(List<Selection<?>> selections) {
-		//noinspection unchecked,rawtypes
-		return array( Object[].class, (List) selections );
+		return arrayInternal( Object[].class,
+				selections.stream().map( selection -> (SqmSelectableNode<?>) selection ).toList() );
 	}
 
 	@Override
-	public <Y> JpaCompoundSelection<Y> array(Class<Y> resultClass, Selection<?>[] selections) {
-		//noinspection unchecked
-		return array( resultClass, (List<SqmSelectableNode<?>>) (List<?>) Arrays.asList( selections ) );
+	public <Y> JpaCompoundSelection<Y> array(Class<Y> resultClass, Selection<?>... selections) {
+		return arrayInternal( resultClass,
+				Arrays.stream( selections ).map( selection -> (SqmSelectableNode<?>) selection ).toList() );
 	}
 
 	@Override
 	public <Y> JpaCompoundSelection<Y> array(Class<Y> resultClass, List<? extends JpaSelection<?>> selections) {
-		//noinspection rawtypes,unchecked
-		checkMultiselect( (List) selections );
+		return arrayInternal( resultClass,
+				selections.stream().map( selection -> (SqmSelectableNode<?>) selection ).toList() );
+	}
+
+	public <Y> JpaCompoundSelection<Y> arrayInternal(Class<Y> resultClass, List<? extends SqmSelectableNode<?>> selections) {
+		checkMultiselect( selections );
 		final JavaType<Y> javaType = getTypeConfiguration().getJavaTypeRegistry().getDescriptor( resultClass );
-		//noinspection unchecked
-		return new SqmJpaCompoundSelection<>( (List<SqmSelectableNode<?>>) selections, javaType, this );
+		return new SqmJpaCompoundSelection<>( selections, javaType, this );
 	}
 
 	@Override
-	public <Y> JpaCompoundSelection<Y> construct(Class<Y> resultClass, Selection<?>[] arguments) {
-		//noinspection unchecked
-		return construct( resultClass, (List<JpaSelection<?>>) (List<?>) Arrays.asList( arguments ) );
+	public <Y> JpaCompoundSelection<Y> construct(Class<Y> resultClass, Selection<?>... arguments) {
+		return constructInternal( resultClass,
+				Arrays.stream( arguments ).map( arg -> (SqmSelectableNode<?>) arg ).toList() );
 	}
 
 	@Override
 	public <Y> JpaCompoundSelection<Y> construct(Class<Y> resultClass, List<? extends JpaSelection<?>> arguments) {
-		//noinspection unchecked,rawtypes
-		checkMultiselect( (List) arguments );
-		final SqmDynamicInstantiation<Y> instantiation;
+		return constructInternal( resultClass,
+				arguments.stream().map( arg -> (SqmSelectableNode<?>) arg ).toList() );
+	}
+
+	private <Y> JpaCompoundSelection<Y> constructInternal(Class<Y> resultClass, List<? extends SqmSelectableNode<?>> arguments) {
+		checkMultiselect( arguments );
+		final SqmDynamicInstantiation<Y> instantiation = createInstantiation( resultClass );
+		for ( SqmSelectableNode<?> argument : arguments ) {
+			final SqmDynamicInstantiationArgument<?> arg =
+					new SqmDynamicInstantiationArgument<>( argument, argument.getAlias(), this );
+			instantiation.addArgument( arg );
+		}
+		return instantiation;
+	}
+
+	@SuppressWarnings("unchecked")
+	private <Y> SqmDynamicInstantiation<Y> createInstantiation(Class<Y> resultClass) {
 		if ( List.class.equals( resultClass ) ) {
-			//noinspection unchecked
-			instantiation = (SqmDynamicInstantiation<Y>) SqmDynamicInstantiation.forListInstantiation( this );
+			return (SqmDynamicInstantiation<Y>) SqmDynamicInstantiation.forListInstantiation( this );
 		}
 		else if ( Map.class.equals( resultClass ) ) {
-			//noinspection unchecked
-			instantiation = (SqmDynamicInstantiation<Y>) SqmDynamicInstantiation.forMapInstantiation( this );
+			return (SqmDynamicInstantiation<Y>) SqmDynamicInstantiation.forMapInstantiation( this );
 		}
 		else {
-			instantiation = SqmDynamicInstantiation.forClassInstantiation( resultClass, this );
+			return SqmDynamicInstantiation.forClassInstantiation( resultClass, this );
 		}
-
-		for ( Selection<?> argument : arguments ) {
-			final SqmSelectableNode<?> arg = (SqmSelectableNode<?>) argument;
-			instantiation.addArgument(
-					new SqmDynamicInstantiationArgument<>( arg, argument.getAlias(), this )
-			);
-		}
-
-		return instantiation;
 	}
 
 	/**
@@ -985,8 +992,8 @@ else if ( Map.class.equals( resultClass ) ) {
 	 * <i>&quot;An argument to the multiselect method must not be a tuple-
 	 * or array-valued compound selection item.&quot;</i>
 	 */
-	void checkMultiselect(List<Selection<?>> selections) {
-		final HashSet<String> aliases = new HashSet<>( CollectionHelper.determineProperSizing( selections.size() ) );
+	void checkMultiselect(List<? extends Selection<?>> selections) {
+		final HashSet<String> aliases = new HashSet<>( determineProperSizing( selections.size() ) );
 
 		for ( Selection<?> it : selections ) {
 			final JpaSelection<?> selection = (JpaSelection<?>) it;

hibernate-core/src/main/java/org/hibernate/query/sqm/sql/BaseSqmToSqlAstConverter.java

@@ -2437,7 +2437,7 @@ private int indexOfExpression(int offset, List<? extends SqmAliasedNode<?>> sele
 				offset = -subResult - i;
 			}
 			else if ( selectableNode instanceof SqmJpaCompoundSelection<?> compoundSelection ) {
-				final List<SqmSelectableNode<?>> selectionItems = compoundSelection.getSelectionItems();
+				final List<? extends SqmSelectableNode<?>> selectionItems = compoundSelection.getSelectionItems();
 				for ( int j = 0; j < selectionItems.size(); j++ ) {
 					if ( selectionItems.get( j ) == node ) {
 						return offset + i + j;

hibernate-core/src/main/java/org/hibernate/query/sqm/tree/from/SqmFromClause.java

@@ -6,14 +6,16 @@
 
 import java.io.Serializable;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 import java.util.function.Consumer;
 
 import org.hibernate.internal.util.collections.CollectionHelper;
 import org.hibernate.query.sqm.tree.SqmCopyContext;
 import org.hibernate.query.sqm.tree.domain.SqmTreatedPath;
 
+import static java.util.Collections.emptyList;
+import static java.util.Collections.unmodifiableList;
+
 /**
  * Contract representing a from clause.
  * <p>
@@ -50,7 +52,7 @@ public SqmFromClause copy(SqmCopyContext context) {
 	 * mutate the roots
 	 */
 	public List<SqmRoot<?>> getRoots() {
-		return domainRoots == null ? Collections.emptyList() : Collections.unmodifiableList( domainRoots );
+		return domainRoots == null ? emptyList() : unmodifiableList( domainRoots );
 	}
 
 	/**