HHH-17483 Fix applyDiscriminator treat for nested inheritance subtypes

Also small fix to joined-inheritance pruning.

hibernate-core/src/main/java/org/hibernate/persister/entity/AbstractEntityPersister.java

@@ -2999,7 +2999,7 @@ public abstract class AbstractEntityPersister
 
 	public abstract Map<Object, String> getSubclassByDiscriminatorValue();
 
-	protected abstract boolean needsDiscriminator();
+	public abstract boolean needsDiscriminator();
 
 	protected boolean isDiscriminatorFormula() {
 		return false;
@@ -3092,7 +3092,25 @@ public abstract class AbstractEntityPersister
 			TableGroup tableGroup,
 			SqlAstCreationState creationState) {
 		if ( needsDiscriminator() ) {
-			pruneForSubclasses( tableGroup, Collections.singletonMap( getEntityName(), EntityNameUse.TREAT ) );
+			assert !creationState.supportsEntityNameUsage() : "Entity name usage should have been used instead";
+			final Map<String, EntityNameUse> entityNameUseMap;
+			final Collection<EntityMappingType> subMappingTypes = getSubMappingTypes();
+			if ( subMappingTypes.isEmpty() ) {
+				entityNameUseMap = Collections.singletonMap( getEntityName(), EntityNameUse.TREAT );
+			}
+			else {
+				entityNameUseMap = new HashMap<>( 1 + subMappingTypes.size() );
+				entityNameUseMap.put( getEntityName(), EntityNameUse.TREAT );
+				// We need to register TREAT uses for all subtypes when pruning
+				for ( EntityMappingType subMappingType : subMappingTypes ) {
+					entityNameUseMap.put( subMappingType.getEntityName(), EntityNameUse.TREAT );
+				}
+				if ( isInherited() ) {
+					// Make sure the table group includes the root table when needed for TREAT
+					tableGroup.resolveTableReference( getRootTableName() );
+				}
+			}
+			pruneForSubclasses( tableGroup, entityNameUseMap );
 		}
 	}
 

hibernate-core/src/main/java/org/hibernate/persister/entity/JoinedSubclassEntityPersister.java

@@ -742,7 +742,7 @@ public class JoinedSubclassEntityPersister extends AbstractEntityPersister {
 	}
 
 	@Override
-	protected boolean needsDiscriminator() {
+	public boolean needsDiscriminator() {
 		return forceDiscriminator;
 	}
 
@@ -1323,15 +1323,17 @@ public class JoinedSubclassEntityPersister extends AbstractEntityPersister {
 				final String[] subclassTableNames = persister.getSubclassTableNames();
 				// Build the intersection of all tables names that are of the class or super class
 				// These are the tables that can be safely inner joined
-				if ( tablesToInnerJoin.isEmpty() ) {
+				final Set<String> classOrSuperclassTables = new HashSet<>( subclassTableNames.length );
-					for ( int i = 0; i < subclassTableNames.length; i++ ) {
+				for ( int i = 0; i < subclassTableNames.length; i++ ) {
-						if ( persister.isClassOrSuperclassTable[i] ) {
+					if ( persister.isClassOrSuperclassTable[i] ) {
-							tablesToInnerJoin.add( subclassTableNames[i] );
+						classOrSuperclassTables.add( subclassTableNames[i] );
-						}
 					}
 				}
+				if ( tablesToInnerJoin.isEmpty() ) {
+					tablesToInnerJoin.addAll( classOrSuperclassTables );
+				}
 				else {
-					tablesToInnerJoin.retainAll( Arrays.asList( subclassTableNames ) );
+					tablesToInnerJoin.retainAll( classOrSuperclassTables );
 				}
 				if ( useKind == EntityNameUse.UseKind.FILTER && explicitDiscriminatorColumnName == null ) {
 					// If there is no discriminator column,

hibernate-core/src/main/java/org/hibernate/persister/entity/SingleTableEntityPersister.java

@@ -548,7 +548,7 @@ public class SingleTableEntityPersister extends AbstractEntityPersister {
 	}
 
 	@Override
-	protected boolean needsDiscriminator() {
+	public boolean needsDiscriminator() {
 		return forceDiscriminator || isInherited();
 	}
 

hibernate-core/src/main/java/org/hibernate/persister/entity/UnionSubclassEntityPersister.java

@@ -286,7 +286,7 @@ public class UnionSubclassEntityPersister extends AbstractEntityPersister {
 	}
 
 	@Override
-	protected boolean needsDiscriminator() {
+	public boolean needsDiscriminator() {
 		return false;
 	}
 

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

@@ -3118,20 +3118,30 @@ public abstract class BaseSqmToSqlAstConverter<T extends Statement> extends Base
 				);
 			}
 		}
-		if ( useKind == EntityNameUse.UseKind.TREAT || useKind == EntityNameUse.UseKind.PROJECTION ) {
+		// If we encounter a treat or projection use, we also want register the use for all subtypes.
-			// If we encounter a treat use, we also want register the use for all subtypes.
+		// We do this here to not have to expand entity name uses during pruning later on
-			// We do this here to not have to expand entity name uses during pruning later on
+		if ( useKind == EntityNameUse.UseKind.TREAT ) {
 			for ( EntityMappingType subType : persister.getSubMappingTypes() ) {
 				entityNameUses.compute(
 						subType.getEntityName(),
 						(s, existingUse) -> finalEntityNameUse.stronger( existingUse )
 				);
-				if ( useKind == EntityNameUse.UseKind.PROJECTION ) {
+			}
-					actualTableGroup.resolveTableReference(
+			if ( persister.isInherited() && persister.needsDiscriminator() ) {
-							null,
+				// Make sure the table group includes the root table when needed for TREAT
-							subType.getEntityPersister().getMappedTableDetails().getTableName()
+				actualTableGroup.resolveTableReference( persister.getRootTableName() );
-					);
+			}
-				}
+		}
+		else if ( useKind == EntityNameUse.UseKind.PROJECTION ) {
+			for ( EntityMappingType subType : persister.getSubMappingTypes() ) {
+				entityNameUses.compute(
+						subType.getEntityName(),
+						(s, existingUse) -> finalEntityNameUse.stronger( existingUse )
+				);
+				actualTableGroup.resolveTableReference(
+						null,
+						subType.getEntityPersister().getMappedTableDetails().getTableName()
+				);
 			}
 		}
 	}
@@ -8238,7 +8248,9 @@ public abstract class BaseSqmToSqlAstConverter<T extends Statement> extends Base
 							if ( entityMappingType.getSuperMappingType() != null ) {
 								// A joined table group was created by an enabled entity graph or fetch profile,
 								// and it's of an inheritance subtype, so we should apply the discriminator
-								entityMappingType.applyDiscriminator( null, null, actualTableGroup, this );
+								getCurrentClauseStack().push( Clause.FROM );
+								registerEntityNameUsage( actualTableGroup, EntityNameUse.TREAT, entityMappingType.getEntityName() );
+								getCurrentClauseStack().pop();
 							}
 						}
 					}