HHH-2862, HHH-1914 Make collection replacement a first class operation on PersistenceContext

Author: beikov

hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java

@@ -1075,6 +1075,35 @@ public void addNewCollection(CollectionPersister persister, PersistentCollection
 		addCollection( collection, persister );
 	}
 
+	@Override
+	public void replaceCollection(CollectionPersister persister, PersistentCollection<?> oldCollection, PersistentCollection<?> collection) {
+		if ( !oldCollection.isDirectlyAccessible() ) {
+			throw new HibernateException(
+					"Replacement of not directly accessible collection found: " + oldCollection.getRole() );
+		}
+		assert !collection.isDirectlyAccessible();
+		final IdentityMap<PersistentCollection<?>, CollectionEntry> collectionEntries = getOrInitializeCollectionEntries();
+		final CollectionEntry oldEntry = collectionEntries.remove( oldCollection );
+		final CollectionEntry entry;
+		if ( oldEntry.getLoadedPersister() != null ) {
+			// This is an already existing/loaded collection so ensure the loadedPersister is initialized
+			entry = new CollectionEntry( collection, session.getFactory() );
+		}
+		else {
+			// A newly wrapped collection
+			entry = new CollectionEntry( persister, collection );
+		}
+		collectionEntries.put( collection, entry );
+		final Object key = collection.getKey();
+		if ( key != null ) {
+			final CollectionKey collectionKey = new CollectionKey( entry.getLoadedPersister(), key );
+			final PersistentCollection<?> old = addCollectionByKey( collectionKey, collection );
+			if ( old == null ) {
+				throw new HibernateException( "No collection for replacement found: " + collectionKey.getRole() );
+			}
+		}
+	}
+
 	/**
 	 * Add a collection to the cache, with a given collection entry.
 	 *

hibernate-core/src/main/java/org/hibernate/engine/spi/PersistenceContext.java

@@ -25,6 +25,7 @@
 
 import org.checkerframework.checker.nullness.qual.Nullable;
 
+
 /**
  * Represents the state of "stuff" Hibernate is tracking, including (not exhaustive):
  * <ul>
@@ -408,6 +409,15 @@ void addInitializedDetachedCollection(
 			CollectionPersister collectionPersister,
 			PersistentCollection<?> collection);
 
+	/**
+	 * Replaces a directly accessible collection with the given one
+	 *
+	 * @param oldCollection
+	 * @param collection The collection to be associated with the persistence context
+	 * @since 7.0
+	 */
+	void replaceCollection(CollectionPersister persister, PersistentCollection<?> oldCollection, PersistentCollection<?> collection);
+
 	/**
 	 * add a collection we just pulled out of the cache (does not need initializing)
 	 */

hibernate-core/src/main/java/org/hibernate/type/CollectionType.java

@@ -661,18 +661,19 @@ private Object replaceOriginal(
 			final boolean wasClean =
 					target instanceof PersistentCollection<?> collection
 							&& !collection.isDirty();
-			if ( target instanceof PersistentCollection<?> existingPersistentCollection
-				&& existingPersistentCollection.isDirectlyAccessible() ) {
+			if ( target instanceof PersistentCollection<?> oldCollection
+				&& oldCollection.isDirectlyAccessible() ) {
 				// When a replace/merge is requested and the underlying collection is directly accessible,
 				// use a new persistent collection, to avoid potential issues
 				// like the underlying collection being unmodifiable and hence failing the element replacement
 				final CollectionPersister collectionPersister = getPersister( session );
-				final Object key = existingPersistentCollection.getKey();
-				final PersistentCollection<?> persistentCollection = instantiate( session, collectionPersister, key );
-				persistentCollection.initializeEmptyCollection( collectionPersister );
-				persistentCollection.setSnapshot( key, existingPersistentCollection.getRole(), existingPersistentCollection.getStoredSnapshot() );
-				session.getPersistenceContextInternal().addInitializedDetachedCollection( collectionPersister, persistentCollection );
-				target = persistentCollection;
+				final Object key = oldCollection.getKey();
+				final PersistentCollection<?> newCollection = instantiate( session, collectionPersister, key );
+				newCollection.initializeEmptyCollection( collectionPersister );
+				newCollection.setSnapshot( key, oldCollection.getRole(), oldCollection.getStoredSnapshot() );
+				session.getPersistenceContextInternal()
+						.replaceCollection( collectionPersister, oldCollection, newCollection );
+				target = newCollection;
 			}
 			//TODO: this is a little inefficient, don't need to do a whole
 			//      deep replaceElements() call