Changes for HR

Author: DavideD

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

@@ -4,26 +4,7 @@
  */
 package org.hibernate.engine.internal;
 
-import java.io.IOException;
-import java.io.InvalidObjectException;
-import java.io.ObjectInputStream;
-import java.io.ObjectOutputStream;
-import java.io.Serializable;
-import java.lang.invoke.MethodHandles;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.IdentityHashMap;
-import java.util.Iterator;
-import java.util.Map;
-import java.util.Map.Entry;
-import java.util.NoSuchElementException;
-import java.util.function.BiConsumer;
-import java.util.function.Consumer;
-import java.util.function.Supplier;
-
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.hibernate.AssertionFailure;
 import org.hibernate.Hibernate;
 import org.hibernate.HibernateException;
@@ -67,10 +48,27 @@
 import org.hibernate.sql.results.jdbc.spi.JdbcValuesSourceProcessingState;
 import org.hibernate.sql.results.spi.LoadContexts;
 import org.hibernate.type.CollectionType;
-
 import org.jboss.logging.Logger;
 
-import org.checkerframework.checker.nullness.qual.Nullable;
+import java.io.IOException;
+import java.io.InvalidObjectException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.Serializable;
+import java.lang.invoke.MethodHandles;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.IdentityHashMap;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.NoSuchElementException;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 import static org.hibernate.engine.internal.ManagedTypeHelper.asHibernateProxy;
 import static org.hibernate.engine.internal.ManagedTypeHelper.asManagedEntity;
@@ -330,10 +328,7 @@ public Object[] getDatabaseSnapshot(Object id, EntityPersister persister) throws
 		}
 		else {
 			final Object[] snapshot = persister.getDatabaseSnapshot( id, session );
-			if ( entitySnapshotsByKey == null ) {
-				entitySnapshotsByKey = CollectionHelper.mapOfSize( INIT_COLL_SIZE );
-			}
-			entitySnapshotsByKey.put( key, snapshot == null ? NO_ROW : snapshot );
+			getOrInitializeEntitySnapshotsByKey().put( key, snapshot == null ? NO_ROW : snapshot );
 			return snapshot;
 		}
 	}
@@ -1187,7 +1182,8 @@ public void initializeNonLazyCollections() throws HibernateException {
 		initializeNonLazyCollections( PersistentCollection::forceInitialization );
 	}
 
-	protected void initializeNonLazyCollections(Consumer<PersistentCollection<?>> initializeAction ) {
+	@Override
+	public void initializeNonLazyCollections(Consumer<PersistentCollection<?>> initializeAction) {
 		if ( loadCounter == 0 ) {
 			LOG.trace( "Initializing non-lazy collections" );
 
@@ -1314,6 +1310,19 @@ public Map<EntityKey,Object> getEntitiesByKey() {
 		return result;
 	}
 
+	@Override
+	public Map<EntityKey, Object> getEntitySnapshotsByKey() {
+		return entitySnapshotsByKey;
+	}
+
+	@Override
+	public Map<EntityKey, Object> getOrInitializeEntitySnapshotsByKey() {
+		if ( entitySnapshotsByKey == null ) {
+			entitySnapshotsByKey = CollectionHelper.mapOfSize( INIT_COLL_SIZE );
+		}
+		return entitySnapshotsByKey;
+	}
+
 	@Override
 	public Map<EntityKey, EntityHolder> getEntityHoldersByKey() {
 		//noinspection unchecked,rawtypes

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

@@ -444,6 +444,8 @@ CollectionEntry addInitializedCollection(
 	 */
 	void initializeNonLazyCollections() throws HibernateException;
 
+	void initializeNonLazyCollections(Consumer<PersistentCollection<?>> initializeAction);
+
 	/**
 	 * Get the {@code PersistentCollection} object for an array
 	 */
@@ -533,6 +535,12 @@ EntityHolder claimEntityHolderIfPossible(
 	@Internal
 	Map<EntityKey,Object> getEntitiesByKey();
 
+	@Internal
+	Map<EntityKey,Object> getEntitySnapshotsByKey();
+
+	@Internal
+	Map<EntityKey,Object> getOrInitializeEntitySnapshotsByKey();
+
 	/**
 	 * Doubly internal
 	 */