Graphpocalypse: major revision/refactoring of EntityGraph support

- important simplification

hibernate-core/src/main/java/org/hibernate/graph/AttributeNode.java

@@ -4,11 +4,12 @@
  */
 package org.hibernate.graph;
 
-import java.util.Map;
 
 import org.hibernate.metamodel.model.domain.ManagedDomainType;
 import org.hibernate.metamodel.model.domain.PersistentAttribute;
 
+import java.util.Map;
+
 /**
  * Extends the JPA-defined {@link AttributeNode} with additional operations.
  *
@@ -26,8 +27,8 @@ public interface AttributeNode<J> extends GraphNode<J>, jakarta.persistence.Attr
 	SubGraph<?> makeSubGraph();
 	SubGraph<?> makeKeySubGraph();
 
-	<S> SubGraph<S> makeSubGraph(Class<S> type);
+	<S> SubGraph<S> makeSubGraph(Class<S> subtype);
-	<S> SubGraph<S> makeKeySubGraph(Class<S> type);
+	<S> SubGraph<S> makeKeySubGraph(Class<S> subtype);
 
 	<S> SubGraph<S> makeSubGraph(ManagedDomainType<S> subtype);
 	<S> SubGraph<S> makeKeySubGraph(ManagedDomainType<S> subtype);

hibernate-core/src/main/java/org/hibernate/graph/EntityGraphs.java

@@ -61,7 +61,6 @@ private static <T> EntityGraph<T> mergeInternal(
 			for ( jakarta.persistence.Graph<T> graph : graphs ) {
 				merged.merge( (GraphImplementor<T>) graph );
 			}
-
 		}
 		return merged;
 	}

hibernate-core/src/main/java/org/hibernate/graph/Graph.java

@@ -11,6 +11,7 @@
 import org.hibernate.metamodel.model.domain.ManagedDomainType;
 import org.hibernate.metamodel.model.domain.MapPersistentAttribute;
 import org.hibernate.metamodel.model.domain.PersistentAttribute;
+import org.hibernate.metamodel.model.domain.PluralPersistentAttribute;
 
 /**
  * A container for {@link AttributeNode} references.
@@ -130,6 +131,8 @@ default <Y> AttributeNode<Y> addAttributeNode(Attribute<? super J, Y> attribute)
 
 	<Y extends J> SubGraph<Y> addTreatedSubGraph(Class<Y> type);
 
+	<Y extends J> SubGraph<Y> addTreatedSubGraph(ManagedDomainType<Y> type);
+
 	/**
 	 * Create and return a new (mutable) {@link SubGraph} associated with
 	 * the named {@link AttributeNode}.
@@ -151,18 +154,22 @@ default <Y> AttributeNode<Y> addAttributeNode(Attribute<? super J, Y> attribute)
 
 	<AJ> SubGraph<AJ> addSubGraph(PersistentAttribute<? super J, ? super AJ> attribute, Class<AJ> type);
 
-	<AJ> SubGraph<AJ> addSubGraph(MapPersistentAttribute<? super J, ? super AJ, ?> attribute, ManagedDomainType<AJ> type);
+	<AJ> SubGraph<AJ> addSubGraph(PersistentAttribute<? super J, ? super AJ> attribute, ManagedDomainType<AJ> type);
+
+	<AJ> SubGraph<AJ> addElementSubGraph(PluralPersistentAttribute<? super J, ?, ? super AJ> attribute, Class<AJ> type);
+
+	<AJ> SubGraph<AJ> addElementSubGraph(PluralPersistentAttribute<? super J, ?, ? super AJ> attribute, ManagedDomainType<AJ> type);
 
 	@Deprecated
 	<AJ> SubGraph<AJ> addKeySubGraph(String attributeName);
 
 	<AJ> SubGraph<AJ> addKeySubGraph(String attributeName, Class<AJ> type);
 
-	<AJ> SubGraph<AJ> addKeySubGraph(PersistentAttribute<? super J, ? super AJ> attribute, ManagedDomainType<AJ> type);
+	<AJ> SubGraph<AJ> addKeySubGraph(MapPersistentAttribute<? super J, ? super AJ, ?> attribute, ManagedDomainType<AJ> type);
 
 	@Override
 	default <Y> SubGraph<Y> addTreatedSubgraph(Attribute<? super J, ? super Y> attribute, Class<Y> type) {
-		return addSubGraph( (PersistentAttribute<? super J, ? super Y>) attribute, type );
+		return addSubGraph( (PersistentAttribute<? super J, ? super Y>) attribute ).addTreatedSubGraph( type );
 	}
 
 	@Override
@@ -172,7 +179,7 @@ default <X> SubGraph<X> addSubgraph(Attribute<? super J, X> attribute) {
 
 	@Override
 	default <X> SubGraph<? extends X> addSubgraph(Attribute<? super J, X> attribute, Class<? extends X> type) {
-		return addSubGraph( (PersistentAttribute<? super J, X>) attribute, type );
+		return addSubGraph( (PersistentAttribute<? super J, X>) attribute ).addTreatedSubGraph( type );
 	}
 
 	@Override
@@ -182,7 +189,7 @@ default <X> SubGraph<X> addSubgraph(String name) {
 
 	@Override
 	default <X> SubGraph<X> addSubgraph(String name, Class<X> type) {
-		return addSubGraph( name, type );
+		return addSubGraph( name ).addTreatedSubGraph( type );
 	}
 
 	@Override
@@ -192,7 +199,7 @@ default <X> SubGraph<X> addKeySubgraph(String name) {
 
 	@Override
 	default <X> SubGraph<X> addKeySubgraph(String name, Class<X> type) {
-		return addKeySubGraph( name, type );
+		return addKeySubGraph( name ).addTreatedSubGraph( type );
 	}
 
 	/**

hibernate-core/src/main/java/org/hibernate/graph/internal/AbstractGraph.java

@@ -23,6 +23,7 @@
 import org.hibernate.metamodel.model.domain.ManagedDomainType;
 import org.hibernate.metamodel.model.domain.MapPersistentAttribute;
 import org.hibernate.metamodel.model.domain.PersistentAttribute;
+import org.hibernate.metamodel.model.domain.PluralPersistentAttribute;
 import org.hibernate.query.sqm.SqmPathSource;
 
 import jakarta.persistence.metamodel.Attribute;
@@ -39,8 +40,8 @@
 public abstract class AbstractGraph<J> extends AbstractGraphNode<J> implements GraphImplementor<J> {
 
 	private final ManagedDomainType<J> managedType;
+	private final Map<Class<?>, SubGraphImplementor<?>> subgraphs = new HashMap<>(1);
 	private Map<PersistentAttribute<? super J,?>, AttributeNodeImplementor<?>> attributeNodes;
-	private List<SubGraphImplementor<? extends J>> subgraphs;
 
 	public AbstractGraph(ManagedDomainType<J> managedType, boolean mutable) {
 		super( mutable );
@@ -50,8 +51,7 @@ public AbstractGraph(ManagedDomainType<J> managedType, boolean mutable) {
 	protected AbstractGraph(ManagedDomainType<J> managedType, GraphImplementor<? super J> graph, boolean mutable) {
 		this( managedType, mutable );
 		attributeNodes = new HashMap<>( graph.getAttributeNodesByAttribute().size() );
-		graph.getAttributeNodesByAttribute()
+		mergeInternal( graph, mutable );
-				.forEach( (attribute, node) -> attributeNodes.put( attribute, node.makeCopy( mutable ) ) );
 	}
 
 	protected AbstractGraph(GraphImplementor<J> graph, boolean mutable) {
@@ -76,42 +76,61 @@ public RootGraphImplementor<J> makeRootGraph(String name, boolean mutable) {
 	}
 
 	@Override
-	public void merge(GraphImplementor<J> graph) {
+	public void merge(GraphImplementor<? super J> graph) {
+		merge( graph, true );
+	}
+
+	@Override
+	public void merge(GraphImplementor<? super J> graph, boolean mutable) {
 		if ( graph != null ) {
 			verifyMutability();
+			mergeInternal( graph, mutable );
+		}
+	}
+
+	private void mergeInternal(GraphImplementor<? super J> graph, boolean mutable) {
 		graph.getAttributeNodesByAttribute().forEach( (attribute, node) -> {
 			final AttributeNodeImplementor<?> existingNode = findAttributeNode( attribute );
 			if ( existingNode != null ) {
 				// keep the local one, but merge in the incoming one
-					mergeNode( node, existingNode );
+				mergeNode( node, existingNode, mutable );
 			}
 			else {
-					addAttributeNode( attribute, node.makeCopy( true ) );
+				addAttributeNode( attribute, node.makeCopy( mutable ), mutable );
+			}
+		} );
+		graph.getSubGraphMap().forEach( (type, subgraph) -> {
+			final SubGraphImplementor<?> existing = subgraphs.get( type );
+			if ( existing != null ) {
+				existing.merge( (SubGraphImplementor) subgraph, mutable );
+			}
+			else {
+				subgraphs.put( type, subgraph.makeCopy( mutable ) );
 			}
 		} );
 	}
-	}
 
-	private static <T> void mergeNode(AttributeNodeImplementor<?> node, AttributeNodeImplementor<T> existingNode) {
+	private static <T> void mergeNode(
+			AttributeNodeImplementor<?> node, AttributeNodeImplementor<T> existingNode, boolean mutable) {
 		if ( existingNode.getAttributeDescriptor() == node.getAttributeDescriptor() ) {
 			@SuppressWarnings("unchecked") // safe, we just checked
 			final AttributeNodeImplementor<T> castNode = (AttributeNodeImplementor<T>) node;
-			existingNode.merge( castNode );
+			existingNode.merge( castNode, mutable );
 		}
 		else {
 			throw new AssertionFailure( "Attributes should have been identical" );
 		}
 	}
 
-	private <T> void addAttributeNode(PersistentAttribute<? super J, ?> attribute, AttributeNodeImplementor<T> node) {
+	private <T> void addAttributeNode(
+			PersistentAttribute<? super J, ?> attribute, AttributeNodeImplementor<T> node, boolean mutable) {
 		final AttributeNodeImplementor<T> attributeNode = getNodeForPut( node.getAttributeDescriptor() );
 		if ( attributeNode == null ) {
 			attributeNodes.put( attribute, node );
 		}
 		else {
 			// we assume the subgraph has been properly copied if needed
-			node.getSubGraphMap().forEach( (subtype, subgraph) -> attributeNode.addSubGraph( subgraph ) );
+			node.merge( attributeNode, mutable );
-			node.getKeySubGraphMap().forEach( (subtype, subgraph) -> attributeNode.addKeySubGraph( subgraph ) );
 		}
 	}
 
@@ -126,8 +145,8 @@ public <AJ> AttributeNodeImplementor<AJ> findAttributeNode(String attributeName)
 		@SuppressWarnings("unchecked") // The JPA API is unsafe by nature
 		final PersistentAttribute<? super J, AJ> persistentAttribute = (PersistentAttribute<? super J, AJ>) attribute;
 		final AttributeNodeImplementor<AJ> node = attribute == null ? null : findAttributeNode( persistentAttribute );
-		if ( node == null && subgraphs != null ) {
+		if ( node == null ) {
-			for ( SubGraphImplementor<? extends J> subgraph : subgraphs ) {
+			for ( SubGraphImplementor<?> subgraph : subgraphs.values() ) {
 				final AttributeNodeImplementor<AJ> subgraphNode = subgraph.findAttributeNode( attributeName );
 				if ( subgraphNode != null ) {
 					return subgraphNode;
@@ -272,12 +291,22 @@ public <AJ> SubGraphImplementor<AJ> addSubGraph(PersistentAttribute<? super J, ?
 	}
 
 	@Override
-	public <AJ> SubGraphImplementor<AJ> addSubGraph(MapPersistentAttribute<? super J, ? super AJ, ?> attribute, ManagedDomainType<AJ> subtype) {
+	public <AJ> SubGraphImplementor<AJ> addSubGraph(PersistentAttribute<? super J, ? super AJ> attribute, ManagedDomainType<AJ> subtype) {
 		return findOrCreateAttributeNode( attribute ).makeSubGraph( subtype );
 	}
 
 	@Override
-	public <AJ> SubGraphImplementor<AJ> addKeySubGraph(PersistentAttribute<? super J, ? super AJ> attribute, ManagedDomainType<AJ> subtype) {
+	public <AJ> SubGraphImplementor<AJ> addElementSubGraph(PluralPersistentAttribute<? super J, ?, ? super AJ> attribute, Class<AJ> type) {
+		return findOrCreateAttributeNode( attribute ).makeSubGraph( type );
+	}
+
+	@Override
+	public <AJ> SubGraphImplementor<AJ> addElementSubGraph(PluralPersistentAttribute<? super J, ?, ? super AJ> attribute, ManagedDomainType<AJ> type) {
+		return findOrCreateAttributeNode( attribute ).makeSubGraph( type );
+	}
+
+	@Override
+	public <AJ> SubGraphImplementor<AJ> addKeySubGraph(MapPersistentAttribute<? super J, ? super AJ, ?> attribute, ManagedDomainType<AJ> subtype) {
 		return findOrCreateAttributeNode( attribute ).makeKeySubGraph( subtype );
 	}
 
@@ -328,13 +357,44 @@ public <E> SubGraphImplementor<E> addTreatedElementSubgraph(
 	}
 
 	@Override
-	public <S extends J> SubGraphImplementor<S> addTreatedSubGraph(Class<S> type) {
+	public <S extends J> SubGraphImplementor<S> addTreatedSubGraph(ManagedDomainType<S> type) {
-		final ManagedDomainType<S> managedDomainType = getGraphedType().getMetamodel().managedType( type );
+		if ( getGraphedType().equals( type ) ) {
-		final SubGraphImpl<S> subgraph = new SubGraphImpl<>( managedDomainType, this, true );
+			//noinspection unchecked
-		if ( subgraphs == null ) {
+			return (SubGraphImplementor<S>) this;
-			subgraphs = new ArrayList<>( 1 );
 		}
-		subgraphs.add( subgraph );
+		else {
+			final Class<S> javaType = type.getJavaType();
+			final SubGraphImplementor<S> castSubgraph = subgraph( javaType );
+			if ( castSubgraph == null ) {
+				final SubGraphImpl<S> subgraph = new SubGraphImpl<>( type, true );
+				subgraphs.put( javaType, subgraph );
 				return subgraph;
 			}
+			else {
+				return castSubgraph;
+			}
+		}
+	}
+
+	private <S extends J> SubGraphImplementor<S> subgraph(Class<S> javaType) {
+		final SubGraphImplementor<?> existing = subgraphs.get( javaType );
+		if ( existing != null ) {
+			@SuppressWarnings("unchecked")
+			final SubGraphImplementor<S> castSubgraph = (SubGraphImplementor<S>) existing;
+			return castSubgraph;
+		}
+		else {
+			return null;
+		}
+	}
+
+	@Override
+	public <S extends J> SubGraphImplementor<S> addTreatedSubGraph(Class<S> type) {
+		return addTreatedSubGraph( getGraphedType().getMetamodel().managedType( type ) );
+	}
+
+	@Override
+	public Map<Class<?>, SubGraphImplementor<?>> getSubGraphMap() {
+		return subgraphs;
+	}
 }

hibernate-core/src/main/java/org/hibernate/graph/internal/AttributeNodeImpl.java

@@ -4,16 +4,19 @@
  */
 package org.hibernate.graph.internal;
 
-import java.util.HashMap;
+import jakarta.persistence.Subgraph;
-import java.util.Map;
-
 import org.hibernate.graph.spi.AttributeNodeImplementor;
 import org.hibernate.graph.spi.SubGraphImplementor;
+import org.hibernate.metamodel.model.domain.DomainType;
 import org.hibernate.metamodel.model.domain.ManagedDomainType;
 import org.hibernate.metamodel.model.domain.PersistentAttribute;
 
+import java.util.HashMap;
+import java.util.Map;
+
 import static java.util.Collections.emptyMap;
 
+
 /**
  * Implementation of {@link jakarta.persistence.AttributeNode}.
  *
@@ -24,25 +27,23 @@ public class AttributeNodeImpl<J>
 		implements AttributeNodeImplementor<J> {
 	private final PersistentAttribute<?, J> attribute;
 
-	private Map<Class<?>, SubGraphImplementor<?>> subgraphMap;
+	private SubGraphImplementor<?> subgraph;
-	private Map<Class<?>, SubGraphImplementor<?>> keySubgraphMap;
+	private SubGraphImplementor<?> keySubgraph;
 
 	public <X> AttributeNodeImpl(PersistentAttribute<X, J> attribute, boolean mutable) {
-		this(attribute, null, null, mutable);
+		this(attribute, mutable, null, null);
 	}
 
 	/**
 	 * Intended only for use from making a copy
 	 */
 	private AttributeNodeImpl(
-			PersistentAttribute<?, J> attribute,
+			PersistentAttribute<?, J> attribute, boolean mutable,
-			Map<Class<?>, SubGraphImplementor<?>> subgraphMap,
+			SubGraphImplementor<?> subgraph, SubGraphImplementor<?> keySubgraph) {
-			Map<Class<?>, SubGraphImplementor<?>> keySubgraphMap,
-			boolean mutable) {
 		super( mutable );
 		this.attribute = attribute;
-		this.subgraphMap = subgraphMap;
+		this.subgraph = subgraph;
-		this.keySubgraphMap = keySubgraphMap;
+		this.keySubgraph = keySubgraph;
 	}
 
 	@Override
@@ -56,191 +57,203 @@ public PersistentAttribute<?, J> getAttributeDescriptor() {
 	}
 
 	@Override
-	public Map<Class<?>, SubGraphImplementor<?>> getSubGraphMap() {
+	public SubGraphImplementor<?> getSubGraph() {
-		return subgraphMap == null ? emptyMap() : subgraphMap;
+		return subgraph;
 	}
 
 	@Override
-	public Map<Class<?>, SubGraphImplementor<?>> getKeySubGraphMap() {
+	public SubGraphImplementor<?> getKeySubGraph() {
-		return keySubgraphMap == null ? emptyMap() : keySubgraphMap;
+		return keySubgraph;
+	}
+
+	private <T> SubGraphImplementor<T> subgraph(ManagedDomainType<T> valueType) {
+		if ( subgraph == null ) {
+			final SubGraphImplementor<T> graph = new SubGraphImpl<>( valueType, true );
+			subgraph = graph;
+			return graph;
+		}
+		else {
+			//noinspection unchecked
+			return (SubGraphImplementor<T>) subgraph;
+		}
 	}
 
 	@Override
 	public SubGraphImplementor<?> makeSubGraph() {
-		return makeSubGraph( (ManagedDomainType<?>) attribute.getValueGraphType() );
+		verifyMutability();
+		return subgraph( (ManagedDomainType<?>) attribute.getValueGraphType() );
 	}
 
 	@Override
-	public <S> SubGraphImplementor<S> makeSubGraph(Class<S> type) {
+	public <S> SubGraphImplementor<S> makeSubGraph(Class<S> subtype) {
-		return makeSubGraph( attribute.getDeclaringType().getMetamodel().managedType( type ) );
+		verifyMutability();
+		final DomainType<?> type = attribute.getValueGraphType();
+		if ( !type.getBindableJavaType().isAssignableFrom( subtype ) ) {
+			throw new IllegalArgumentException( "Not a subtype: " + subtype.getName() );
+		}
+		@SuppressWarnings("unchecked")
+		final ManagedDomainType<? super S> valueType = (ManagedDomainType<? super S>) type;
+		SubGraphImplementor<? super S> subgraph = subgraph( valueType );
+		if ( type.getBindableJavaType() == subtype ) {
+			//noinspection unchecked
+			return (SubGraphImplementor<S>) subgraph;
+		}
+		else {
+			return subgraph.addTreatedSubGraph( subtype );
+		}
 	}
 
 	@Override
 	public <S> SubGraphImplementor<S> makeSubGraph(ManagedDomainType<S> subtype) {
 		verifyMutability();
-		assert subtype != null;
+		final DomainType<?> type = attribute.getValueGraphType();
-		final Class<S> javaType = subtype.getJavaType();
+		final Class<S> javaType = subtype.getBindableJavaType();
-		if ( !attribute.getValueGraphType().getBindableJavaType().isAssignableFrom( javaType ) ) {
+		if ( !type.getBindableJavaType().isAssignableFrom( javaType ) ) {
 			throw new IllegalArgumentException( "Not a subtype: " + javaType.getName() );
 		}
-		final SubGraphImplementor<S> existing = subgraphMap == null ? null : getSubgraph( javaType );
+		@SuppressWarnings("unchecked")
-		if ( existing != null ) {
+		final ManagedDomainType<? super S> valueType = (ManagedDomainType<? super S>) type;
-			return existing;
+		SubGraphImplementor<? super S> subgraph = subgraph( valueType );
+		if ( type.getBindableJavaType() == javaType ) {
+			//noinspection unchecked
+			return (SubGraphImplementor<S>) subgraph;
 		}
 		else {
-			final SubGraphImplementor<S> subGraph = new SubGraphImpl<>( subtype, true );
+			return subgraph.addTreatedSubGraph( subtype );
-			addSubGraph( subGraph );
-			return subGraph;
 		}
 	}
 
-	@Override
+	private <T> SubGraphImplementor<T> keySubgraph(ManagedDomainType<T> keyType) {
-	public void addSubGraph(SubGraphImplementor<?> subgraph) {
+		if ( keySubgraph == null ) {
-		addSubgraph( subgraph );
+			final SubGraphImplementor<T> graph = new SubGraphImpl<>( keyType, true );
-	}
+			keySubgraph = graph;
-
+			return graph;
-	private <T> void addSubgraph(SubGraphImplementor<T> subgraph) {
-		if ( subgraphMap == null ) {
-			subgraphMap = new HashMap<>();
-			subgraphMap.put( subgraph.getClassType(), subgraph );
 		}
 		else {
-			final SubGraphImplementor<T> existing = getSubgraph( subgraph.getClassType() );
+			//noinspection unchecked
-			if ( existing == null ) {
+			return (SubGraphImplementor<T>) keySubgraph;
-				subgraphMap.put( subgraph.getClassType(), subgraph );
-			}
-			else {
-				existing.merge( subgraph );
-			}
-		}
-	}
-
-	@Override
-	public void addKeySubGraph(SubGraphImplementor<?> subgraph) {
-		addKeySubgraph( subgraph );
-	}
-
-	private <T> void addKeySubgraph(SubGraphImplementor<T> subgraph) {
-		if ( keySubgraphMap == null ) {
-			keySubgraphMap = new HashMap<>();
-			keySubgraphMap.put( subgraph.getClassType(), subgraph );
-		}
-		else {
-			final SubGraphImplementor<T> existing = getKeySubgraph( subgraph.getClassType() );
-			if ( existing == null ) {
-				keySubgraphMap.put( subgraph.getClassType(), subgraph );
-			}
-			else {
-				existing.merge( subgraph );
-			}
 		}
 	}
 
 	@Override
 	public SubGraphImplementor<?> makeKeySubGraph() {
-		return makeKeySubGraph( (ManagedDomainType<?>) attribute.getKeyGraphType() );
+		verifyMutability();
+		return keySubgraph( (ManagedDomainType<?>) attribute.getKeyGraphType() );
 	}
 
 	@Override
-	public <S> SubGraphImplementor<S> makeKeySubGraph(Class<S> type) {
+	public <S> SubGraphImplementor<S> makeKeySubGraph(Class<S> subtype) {
-		return makeKeySubGraph( attribute.getDeclaringType().getMetamodel().managedType( type ) );
+		verifyMutability();
+		final DomainType<?> type = attribute.getKeyGraphType();
+		if ( !type.getBindableJavaType().isAssignableFrom( subtype ) ) {
+			throw new IllegalArgumentException( "Not a key subtype: " + subtype.getName() );
+		}
+		@SuppressWarnings("unchecked")
+		final ManagedDomainType<? super S> keyType = (ManagedDomainType<? super S>) type;
+		final SubGraphImplementor<? super S> keySubgraph = keySubgraph( keyType );
+		if ( type.getBindableJavaType() == subtype ) {
+			//noinspection unchecked
+			return (SubGraphImplementor<S>) keySubgraph;
+		}
+		else {
+			return keySubgraph.addTreatedSubGraph( subtype );
+		}
 	}
 
 	@Override
 	public <S> SubGraphImplementor<S> makeKeySubGraph(ManagedDomainType<S> subtype) {
 		verifyMutability();
-		assert subtype != null;
+		final DomainType<?> type = attribute.getKeyGraphType();
-		final Class<S> javaType = subtype.getJavaType();
+		final Class<S> javaType = subtype.getBindableJavaType();
-		if ( !attribute.getKeyGraphType().getBindableJavaType().isAssignableFrom( javaType ) ) {
+		if ( !type.getBindableJavaType().isAssignableFrom( javaType ) ) {
 			throw new IllegalArgumentException( "Not a key subtype: " + javaType.getName() );
 		}
-		final SubGraphImplementor<S> existing = keySubgraphMap == null ? null : getKeySubgraph( javaType );
+		@SuppressWarnings("unchecked")
-		if ( existing != null ) {
+		final ManagedDomainType<? super S> keyType = (ManagedDomainType<? super S>) type;
-			return existing;
+		final SubGraphImplementor<? super S> keySubgraph = keySubgraph( keyType );
+		if ( type.getBindableJavaType() == javaType ) {
+			//noinspection unchecked
+			return (SubGraphImplementor<S>) keySubgraph;
 		}
 		else {
-			final SubGraphImplementor<S> subgraph = new SubGraphImpl<>( subtype, true );
+			return keySubgraph.addTreatedSubGraph( subtype );
-			addKeySubGraph( subgraph );
-			return subgraph;
 		}
 	}
 
 	@Override
 	public AttributeNodeImplementor<J> makeCopy(boolean mutable) {
-		return new AttributeNodeImpl<>(
+		return new AttributeNodeImpl<>( this.attribute, mutable,
-				this.attribute,
+				subgraph == null ? null : subgraph.makeCopy( mutable ),
-				makeMapCopy( mutable, subgraphMap ),
+				keySubgraph == null ? null : keySubgraph.makeCopy( mutable ) );
-				makeMapCopy( mutable, keySubgraphMap ),
-				mutable
-		);
 	}
 
-	private <U,V> Map<Class<? extends U>, SubGraphImplementor<? extends V>> makeMapCopy(
+	@Override
-			boolean mutable,
+	public void merge(AttributeNodeImplementor<J> other, boolean mutable) {
-			Map<Class<? extends U>, SubGraphImplementor<? extends V>> nodeMap) {
+		final SubGraphImplementor<?> otherSubgraph = other.getSubGraph();
-		if ( nodeMap == null ) {
+		if ( otherSubgraph != null ) {
-			return null;
+			if ( subgraph == null ) {
+				subgraph = otherSubgraph.makeCopy( mutable );
 			}
 			else {
-			final HashMap<Class<? extends U>, SubGraphImplementor<? extends V>> map = new HashMap<>( nodeMap.size() );
+				subgraph.merge( (SubGraphImplementor) otherSubgraph, mutable );
-			nodeMap.forEach( (attribute, subgraph) -> map.put( attribute, subgraph.makeCopy( mutable ) ) );
+			}
+		}
+		final SubGraphImplementor<?> otherKeySubgraph = other.getKeySubGraph();
+		if ( otherKeySubgraph != null ) {
+			if ( keySubgraph == null ) {
+				keySubgraph = otherKeySubgraph.makeCopy( mutable );
+			}
+			else {
+				keySubgraph.merge( (SubGraphImplementor) otherKeySubgraph, mutable );
+			}
+		}
+	}
+
+	@Override
+	public Map<Class<?>, SubGraphImplementor<?>> getSubGraphMap() {
+		if ( subgraph == null ) {
+			return emptyMap();
+		}
+		else {
+			final HashMap<Class<?>, SubGraphImplementor<?>> map = new HashMap<>( subgraph.getSubGraphMap() );
+			map.put( attribute.getValueGraphType().getBindableJavaType(), subgraph );
 			return map;
 		}
 	}
 
 	@Override
-	public void merge(AttributeNodeImplementor<J> other) {
+	public Map<Class<?>, SubGraphImplementor<?>> getKeySubGraphMap() {
-		other.getSubGraphMap().values().forEach( this::mergeToSubgraph );
+		if ( keySubgraph == null ) {
-		other.getKeySubGraphMap().values().forEach( this::mergeToKeySubgraph );
+			return emptyMap();
-	}
-
-	private <T> void mergeToKeySubgraph(SubGraphImplementor<T> subgraph) {
-		final SubGraphImplementor<T> existing = getKeySubgraphForPut( subgraph );
-		if ( existing != null ) {
-			existing.merge( subgraph );
 		}
 		else {
-			addKeySubGraph( subgraph.makeCopy( true ) );
+			final HashMap<Class<?>, SubGraphImplementor<?>> map = new HashMap<>( keySubgraph.getSubGraphMap() );
+			map.put( attribute.getKeyGraphType().getJavaType(), keySubgraph );
+			return map;
 		}
 	}
 
-	private <T> void mergeToSubgraph(SubGraphImplementor<T> subgraph) {
+	@Override
-		final SubGraphImplementor<T> existing = getSubgraphForPut( subgraph );
+	public @SuppressWarnings("rawtypes") Map<Class, Subgraph> getSubgraphs() {
-		if ( existing != null ) {
+		if ( subgraph == null ) {
-			existing.merge( subgraph );
+			return emptyMap();
 		}
 		else {
-			addSubGraph( subgraph.makeCopy( true ) );
+			final HashMap<Class, Subgraph> map = new HashMap<>( subgraph.getSubGraphMap() );
+			map.put( attribute.getValueGraphType().getBindableJavaType(), subgraph );
+			return map;
 		}
 	}
 
-	private <T> SubGraphImplementor<T> getSubgraphForPut(SubGraphImplementor<T> subgraph) {
+	@Override
-		if ( subgraphMap == null ) {
+	public @SuppressWarnings("rawtypes") Map<Class, Subgraph> getKeySubgraphs() {
-			subgraphMap = new HashMap<>();
+		if ( keySubgraph == null ) {
-			return null;
+			return emptyMap();
 		}
 		else {
-			return getSubgraph( subgraph.getClassType() );
+			final HashMap<Class, Subgraph> map = new HashMap<>( keySubgraph.getSubGraphMap() );
+			map.put( attribute.getKeyGraphType().getJavaType(), keySubgraph );
+			return map;
 		}
 	}
-
-	private <T> SubGraphImplementor<T> getKeySubgraphForPut(SubGraphImplementor<T> subgraph) {
-		if ( keySubgraphMap == null ) {
-			keySubgraphMap = new HashMap<>();
-			return null;
-		}
-		else {
-			return getKeySubgraph( subgraph.getClassType() );
-		}
-	}
-
-	@SuppressWarnings("unchecked")
-	private <T> SubGraphImplementor<T> getSubgraph(Class<T> incomingSubtype) {
-		return (SubGraphImplementor<T>) subgraphMap.get( incomingSubtype );
-	}
-
-	@SuppressWarnings("unchecked")
-	private <T> SubGraphImplementor<T> getKeySubgraph(Class<T> incomingSubtype) {
-		return (SubGraphImplementor<T>) keySubgraphMap.get( incomingSubtype );
-	}
 }

hibernate-core/src/main/java/org/hibernate/graph/internal/SubGraphImpl.java

@@ -22,10 +22,6 @@ public SubGraphImpl(AbstractGraph<J> original, boolean mutable) {
 		super( original, mutable );
 	}
 
-	public SubGraphImpl(ManagedDomainType<J> managedDomainType, AbstractGraph<? super J> original, boolean mutable) {
-		super( managedDomainType, original, mutable );
-	}
-
 	@Override
 	public SubGraphImplementor<J> makeCopy(boolean mutable) {
 		return new SubGraphImpl<>( this, mutable );

hibernate-core/src/main/java/org/hibernate/graph/spi/AttributeNodeImplementor.java

@@ -4,15 +4,11 @@
  */
 package org.hibernate.graph.spi;
 
-import java.util.Map;
-
-import jakarta.persistence.Subgraph;
-
 import org.hibernate.graph.AttributeNode;
 import org.hibernate.graph.SubGraph;
 import org.hibernate.metamodel.model.domain.ManagedDomainType;
 
-import static java.util.Collections.unmodifiableMap;
+import java.util.Map;
 
 /**
  * Integration version of the {@link AttributeNode} contract
@@ -22,29 +18,6 @@
  */
 public interface AttributeNodeImplementor<J> extends AttributeNode<J>, GraphNodeImplementor<J> {
 
-	Map<Class<?>, SubGraphImplementor<?>> getSubGraphMap();
-	Map<Class<?>, SubGraphImplementor<?>> getKeySubGraphMap();
-
-	@Override
-	default Map<Class<?>, ? extends SubGraph<?>> getSubGraphs() {
-		return unmodifiableMap( getSubGraphMap() );
-	}
-
-	@Override
-	default Map<Class<?>, ? extends SubGraph<?>> getKeySubGraphs() {
-		return unmodifiableMap( getKeySubGraphMap() );
-	}
-
-	@Override // JPA API uses raw types
-	default @SuppressWarnings("rawtypes") Map<Class, Subgraph> getSubgraphs() {
-		return unmodifiableMap( getSubGraphMap() );
-	}
-
-	@Override // JPA API uses raw types
-	default @SuppressWarnings("rawtypes") Map<Class, Subgraph> getKeySubgraphs() {
-		return unmodifiableMap( getKeySubGraphMap() );
-	}
-
 	@Override
 	AttributeNodeImplementor<J> makeCopy(boolean mutable);
 
@@ -55,10 +28,10 @@ public interface AttributeNodeImplementor<J> extends AttributeNode<J>, GraphNode
 	SubGraphImplementor<?> makeKeySubGraph();
 
 	@Override
-	<S> SubGraphImplementor<S> makeSubGraph(Class<S> type);
+	<S> SubGraphImplementor<S> makeSubGraph(Class<S> subtype);
 
 	@Override
-	<S> SubGraphImplementor<S> makeKeySubGraph(Class<S> type);
+	<S> SubGraphImplementor<S> makeKeySubGraph(Class<S> subtype);
 
 	@Override
 	<S> SubGraphImplementor<S> makeSubGraph(ManagedDomainType<S> subtype);
@@ -66,9 +39,23 @@ public interface AttributeNodeImplementor<J> extends AttributeNode<J>, GraphNode
 	@Override
 	<S> SubGraphImplementor<S> makeKeySubGraph(ManagedDomainType<S> subtype);
 
-	void merge(AttributeNodeImplementor<J> other);
+	void merge(AttributeNodeImplementor<J> other, boolean mutable);
 
-	void addSubGraph(SubGraphImplementor<?> subgraph);
+	SubGraphImplementor<?> getSubGraph();
 
-	void addKeySubGraph(SubGraphImplementor<?> subgraph);
+	SubGraphImplementor<?> getKeySubGraph();
+
+	Map<Class<?>, SubGraphImplementor<?>> getSubGraphMap();
+
+	Map<Class<?>, SubGraphImplementor<?>> getKeySubGraphMap();
+
+	@Override
+	default Map<Class<?>, ? extends SubGraph<?>> getSubGraphs() {
+		return getSubGraphMap();
+	}
+
+	@Override
+	default Map<Class<?>, ? extends SubGraph<?>> getKeySubGraphs() {
+		return getKeySubGraphMap();
+	}
 }

hibernate-core/src/main/java/org/hibernate/graph/spi/GraphImplementor.java

@@ -12,6 +12,7 @@
 import org.hibernate.metamodel.model.domain.ManagedDomainType;
 import org.hibernate.metamodel.model.domain.MapPersistentAttribute;
 import org.hibernate.metamodel.model.domain.PersistentAttribute;
+import org.hibernate.metamodel.model.domain.PluralPersistentAttribute;
 
 
 /**
@@ -23,7 +24,9 @@
  */
 public interface GraphImplementor<J> extends Graph<J>, GraphNodeImplementor<J> {
 
-	void merge(GraphImplementor<J> other);
+	void merge(GraphImplementor<? super J> other);
+
+	void merge(GraphImplementor<? super J> other, boolean mutable);
 
 	@Override @Deprecated(forRemoval = true)
 	RootGraphImplementor<J> makeRootGraph(String name, boolean mutable)
@@ -72,11 +75,19 @@ RootGraphImplementor<J> makeRootGraph(String name, boolean mutable)
 	<AJ> SubGraphImplementor<AJ> addKeySubGraph(String attributeName, Class<AJ> subtype);
 
 	@Override
-	<AJ> SubGraphImplementor<AJ> addSubGraph(MapPersistentAttribute<? super J, ? super AJ, ?> attribute, ManagedDomainType<AJ> subtype);
+	<AJ> SubGraphImplementor<AJ> addSubGraph(PersistentAttribute<? super J, ? super AJ> attribute, ManagedDomainType<AJ> subtype);
+
+	<AJ> SubGraphImplementor<AJ> addElementSubGraph(PluralPersistentAttribute<? super J, ?, ? super AJ> attribute, Class<AJ> type);
+
+	<AJ> SubGraphImplementor<AJ> addElementSubGraph(PluralPersistentAttribute<? super J, ?, ? super AJ> attribute, ManagedDomainType<AJ> type);
 
 	@Override
-	<AJ> SubGraphImplementor<AJ> addKeySubGraph(PersistentAttribute<? super J, ? super AJ> attribute, ManagedDomainType<AJ> subtype);
+	<AJ> SubGraphImplementor<AJ> addKeySubGraph(MapPersistentAttribute<? super J, ? super AJ, ?> attribute, ManagedDomainType<AJ> subtype);
 
 	@Override
 	<Y extends J> SubGraphImplementor<Y> addTreatedSubGraph(Class<Y> type);
+
+	<Y extends J> SubGraphImplementor<Y> addTreatedSubGraph(ManagedDomainType<Y> type);
+
+	Map<Class<?>, SubGraphImplementor<?>> getSubGraphMap();
 }

hibernate-core/src/test/java/org/hibernate/orm/test/graph/EntityGraphsTest.java

@@ -25,7 +25,7 @@ private <T> void checkMerge(Class<T> rootType, EntityGraph<T> expected, EntityGr
 	}
 
 	@SafeVarargs
-	private final void checkMerge(EntityGraph<GraphParsingTestEntity> expected, EntityGraph<GraphParsingTestEntity>... graphs) {
+	private void checkMerge(EntityGraph<GraphParsingTestEntity> expected, EntityGraph<GraphParsingTestEntity>... graphs) {
 		checkMerge( GraphParsingTestEntity.class, expected, graphs );
 	}