HHH-19756 Improve sub-part resolution when specifying a treat type

Author: mbellade

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

@@ -5660,49 +5660,30 @@ public ModelPart findSubPart(String name, EntityMappingType treatTargetType) {
 			}
 		}
 
-		if ( treatTargetType != null ) {
-			if ( ! treatTargetType.isTypeOrSuperType( this ) ) {
+		if ( treatTargetType == null ) {
+			final var subDefinedAttribute = findSubPartInSubclassMappings( name );
+			if ( subDefinedAttribute != null ) {
+				return subDefinedAttribute;
+			}
+		}
+		else if ( treatTargetType != this ) {
+			if ( !treatTargetType.isTypeOrSuperType( this ) ) {
 				return null;
 			}
-
-			if ( subclassMappingTypes != null && !subclassMappingTypes.isEmpty() ) {
-				for ( var subMappingType : subclassMappingTypes.values() ) {
-					if ( treatTargetType.isTypeOrSuperType( subMappingType ) ) {
-						final var subDefinedAttribute =
-								subMappingType.findSubTypesSubPart( name, treatTargetType );
-						if ( subDefinedAttribute != null ) {
-							return subDefinedAttribute;
-						}
-					}
-				}
+			// Prefer attributes defined in the treat target type or its subtypes
+			final var treatTypeSubPart = treatTargetType.findSubTypesSubPart( name, null );
+			if ( treatTypeSubPart != null ) {
+				return treatTypeSubPart;
 			}
-		}
-		else {
-			if ( subclassMappingTypes != null && !subclassMappingTypes.isEmpty() ) {
-				ModelPart attribute = null;
-				for ( EntityMappingType subMappingType : subclassMappingTypes.values() ) {
-					final var subDefinedAttribute =
-							subMappingType.findSubTypesSubPart( name, treatTargetType );
-					if ( subDefinedAttribute != null ) {
-						if ( attribute != null && !isCompatibleModelPart( attribute, subDefinedAttribute ) ) {
-							throw new PathException(
-									String.format(
-											Locale.ROOT,
-											"Could not resolve attribute '%s' of '%s' due to the attribute being declared in multiple subtypes '%s' and '%s'",
-											name,
-											getJavaType().getTypeName(),
-											attribute.asAttributeMapping().getDeclaringType()
-													.getJavaType().getTypeName(),
-											subDefinedAttribute.asAttributeMapping().getDeclaringType()
-													.getJavaType().getTypeName()
-									)
-							);
-						}
-						attribute = subDefinedAttribute;
+			else {
+				// If not found, look in the treat target type's supertypes
+				EntityMappingType superType = treatTargetType.getSuperMappingType();
+				while ( superType != this ) {
+					final var superTypeSubPart = superType.findDeclaredAttributeMapping( name );
+					if ( superTypeSubPart != null ) {
+						return superTypeSubPart;
 					}
-				}
-				if ( attribute != null ) {
-					return attribute;
+					superType = superType.getSuperMappingType();
 				}
 			}
 		}
@@ -5725,23 +5706,37 @@ public ModelPart findSubPart(String name, EntityMappingType treatTargetType) {
 		}
 	}
 
+	private ModelPart findSubPartInSubclassMappings(String name) {
+		ModelPart attribute = null;
+		if ( isNotEmpty( subclassMappingTypes ) ) {
+			for ( var subMappingType : subclassMappingTypes.values() ) {
+				final var subDefinedAttribute = subMappingType.findSubTypesSubPart( name, null );
+				if ( subDefinedAttribute != null ) {
+					if ( attribute != null && !isCompatibleModelPart( attribute, subDefinedAttribute ) ) {
+						throw new PathException( String.format(
+								Locale.ROOT,
+								"Could not resolve attribute '%s' of '%s' due to the attribute being declared in multiple subtypes '%s' and '%s'",
+								name,
+								getJavaType().getTypeName(),
+								attribute.asAttributeMapping().getDeclaringType().getJavaType().getTypeName(),
+								subDefinedAttribute.asAttributeMapping().getDeclaringType().getJavaType().getTypeName()
+						) );
+					}
+					attribute = subDefinedAttribute;
+				}
+			}
+		}
+		return attribute;
+	}
+
 	@Override
 	public ModelPart findSubTypesSubPart(String name, EntityMappingType treatTargetType) {
 		final var declaredAttribute = declaredAttributeMappings.get( name );
 		if ( declaredAttribute != null ) {
 			return declaredAttribute;
 		}
 		else {
-			if ( subclassMappingTypes != null && !subclassMappingTypes.isEmpty() ) {
-				for ( var subMappingType : subclassMappingTypes.values() ) {
-					final var subDefinedAttribute =
-							subMappingType.findSubTypesSubPart( name, treatTargetType );
-					if ( subDefinedAttribute != null ) {
-						return subDefinedAttribute;
-					}
-				}
-			}
-			return null;
+			return findSubPartInSubclassMappings( name );
 		}
 	}