HHH-19805 Allow nested entity initializers to claim entity holder

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

@@ -404,29 +404,20 @@ public EntityHolder claimEntityHolderIfPossible(
 		EntityHolderImpl holder = getOrInitializeNewHolder().withEntity( key, key.getPersister(), entity );
 		final EntityHolderImpl oldHolder = entityHolderMap.putIfAbsent( key, newEntityHolder );
 		if ( oldHolder != null ) {
+			// An initializer can't claim an entity holder if it's already initialized
+			if ( oldHolder.isInitialized() ) {
+				return oldHolder;
+			}
 			if ( entity != null ) {
 				assert oldHolder.entity == null || oldHolder.entity == entity;
 				oldHolder.entity = entity;
 			}
-			// Skip setting a new entity initializer if there already is one owner
-			// Also skip if an entity exists which is different from the effective optional object.
-			// The effective optional object is the current object to be refreshed,
-			// which always needs re-initialization, even if already initialized
-			if ( oldHolder.entityInitializer != null
-					|| oldHolder.entity != null && oldHolder.state != EntityHolderState.ENHANCED_PROXY && (
-					processingState.getProcessingOptions().getEffectiveOptionalObject() == null
-							|| oldHolder.entity != processingState.getProcessingOptions().getEffectiveOptionalObject() )
-			) {
-				return oldHolder;
-			}
 			holder = oldHolder;
 		}
 		else {
 			newEntityHolder = null;
 		}
-		assert holder.entityInitializer == null || holder.entityInitializer == initializer;
 		holder.entityInitializer = initializer;
-		processingState.registerLoadingEntityHolder( holder );
 		return holder;
 	}
 

hibernate-core/src/main/java/org/hibernate/internal/util/ImmutableBitSet.java

@@ -4,6 +4,8 @@
  */
 package org.hibernate.internal.util;
 
+import org.checkerframework.checker.nullness.qual.Nullable;
+
 import java.util.Arrays;
 import java.util.BitSet;
 
@@ -26,6 +28,10 @@ private ImmutableBitSet(long[] words) {
 		this.words = words;
 	}
 
+	public static ImmutableBitSet valueOfOrEmpty(@Nullable BitSet bitSet) {
+		return bitSet == null ? EMPTY : valueOf( bitSet );
+	}
+
 	public static ImmutableBitSet valueOf(BitSet bitSet) {
 		final long[] words = bitSet.toLongArray();
 		return words.length == 0 ? EMPTY : new ImmutableBitSet( words );
@@ -70,4 +76,7 @@ public boolean equals(Object obj) {
 		return Arrays.equals( words, set.words );
 	}
 
+	public BitSet toBitSet() {
+		return BitSet.valueOf( words );
+	}
 }

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/internal/PluralAttributeMappingImpl.java

@@ -598,7 +598,7 @@ private Fetch createDelayedCollectionFetch(
 		// returning a domain result assembler that returns LazyPropertyInitializer.UNFETCHED_PROPERTY
 		final EntityMappingType containingEntityMapping = findContainingEntityMapping();
 		final boolean unfetched;
-		if ( fetchParent.getReferencedModePart() == containingEntityMapping
+		if ( fetchParent.getReferencedMappingContainer() == containingEntityMapping
 				&& containingEntityMapping.getEntityPersister().getPropertyLaziness()[getStateArrayPosition()] ) {
 			collectionKeyDomainResult = null;
 			unfetched = true;

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

@@ -84,6 +84,7 @@
 import org.hibernate.internal.CoreMessageLogger;
 import org.hibernate.internal.FilterAliasGenerator;
 import org.hibernate.internal.FilterHelper;
+import org.hibernate.internal.util.ImmutableBitSet;
 import org.hibernate.internal.util.IndexedConsumer;
 import org.hibernate.internal.util.MarkerObject;
 import org.hibernate.internal.util.StringHelper;
@@ -1734,14 +1735,21 @@ protected boolean initializeLazyProperty(
 			final EntityEntry entry,
 			final int index,
 			final Object propValue) {
-		setPropertyValue( entity, lazyPropertyNumbers[index], propValue );
-		if ( entry.getLoadedState() != null ) {
+		final int propertyNumber = lazyPropertyNumbers[index];
+		setPropertyValue( entity, propertyNumber, propValue );
+		if ( entry.getMaybeLazySet() != null ) {
+			var bitSet = entry.getMaybeLazySet().toBitSet();
+			bitSet.set( propertyNumber );
+			entry.setMaybeLazySet( ImmutableBitSet.valueOf( bitSet ) );
+		}
+		final var loadedState = entry.getLoadedState();
+		if ( loadedState != null ) {
 			// object have been loaded with setReadOnly(true); HHH-2236
-			entry.getLoadedState()[lazyPropertyNumbers[index]] = lazyPropertyTypes[index].deepCopy( propValue, factory );
+			loadedState[propertyNumber] = lazyPropertyTypes[index].deepCopy( propValue, factory );
 		}
 		// If the entity has deleted state, then update that as well
 		if ( entry.getDeletedState() != null ) {
-			entry.getDeletedState()[lazyPropertyNumbers[index]] = lazyPropertyTypes[index].deepCopy( propValue, factory );
+			entry.getDeletedState()[propertyNumber] = lazyPropertyTypes[index].deepCopy( propValue, factory );
 		}
 		return fieldName.equals( lazyPropertyNames[index] );
 	}
@@ -1763,6 +1771,11 @@ protected boolean initializeLazyProperty(
 	// Used by Hibernate Reactive
 	protected void initializeLazyProperty(Object entity, EntityEntry entry, Object propValue, int index, Type type) {
 		setPropertyValue( entity, index, propValue );
+		if ( entry.getMaybeLazySet() != null ) {
+			var bitSet = entry.getMaybeLazySet().toBitSet();
+			bitSet.set( index );
+			entry.setMaybeLazySet( ImmutableBitSet.valueOf( bitSet ) );
+		}
 		final Object[] loadedState = entry.getLoadedState();
 		if ( loadedState != null ) {
 			// object have been loaded with setReadOnly(true); HHH-2236