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;
@@ -1598,7 +1599,6 @@ private Object initLazyProperties(
 
 		final var interceptor = asPersistentAttributeInterceptable( entity ).$$_hibernate_getInterceptor();
 		assert interceptor != null : "Expecting bytecode interceptor to be non-null";
-		final Set<String> initializedLazyAttributeNames = interceptor.getInitializedLazyAttributeNames();
 
 		final var lazyAttributesMetadata = getBytecodeEnhancementMetadata().getLazyAttributesMetadata();
 		final String fetchGroup = lazyAttributesMetadata.getFetchGroupName( fieldName );
@@ -1609,29 +1609,32 @@ private Object initLazyProperties(
 		try {
 			Object finalResult = null;
 			final Object[] results = lazySelect.load( id, session );
+			final Set<String> initializedLazyAttributeNames = interceptor.getInitializedLazyAttributeNames();
 			int i = 0;
 			for ( var fetchGroupAttributeDescriptor : fetchGroupAttributeDescriptors ) {
 				final String attributeName = fetchGroupAttributeDescriptor.getName();
-				final boolean previousInitialized = initializedLazyAttributeNames.contains( attributeName );
-				if ( previousInitialized ) {
-					// it's already been initialized (e.g. by a write) so we don't want to overwrite
-					i++;
-					// TODO: we should consider un-marking an attribute as dirty based on the selected value
-					// - we know the current value:
-					//   getPropertyValue( entity, fetchGroupAttributeDescriptor.getAttributeIndex() );
-					// - we know the selected value (see selectedValue below)
-					// - we can use the attribute Type to tell us if they are the same
-					// - assuming entity is a SelfDirtinessTracker we can also know if the attribute is currently
-					//   considered dirty, and if really not dirty we would do the un-marking
-					// - of course that would mean a new method on SelfDirtinessTracker to allow un-marking
+				if ( fieldName.equals( attributeName ) ) {
+					finalResult = results[i];
 				}
-				else {
-					final Object result = results[i++];
-					if ( initializeLazyProperty( fieldName, entity, entry, fetchGroupAttributeDescriptor, result ) ) {
-						finalResult = result;
-						interceptor.attributeInitialized( attributeName );
-					}
+				if ( !initializedLazyAttributeNames.contains( attributeName ) ) {
+					initializeLazyProperty(
+							entity,
+							entry,
+							results[i],
+							getPropertyIndex( attributeName ),
+							fetchGroupAttributeDescriptor.getType()
+					);
 				}
+				// if the attribute has already been initialized (e.g. by a write) we don't want to overwrite
+				i++;
+				// TODO: we should consider un-marking an attribute as dirty based on the selected value
+				// - we know the current value:
+				//   getPropertyValue( entity, fetchGroupAttributeDescriptor.getAttributeIndex() );
+				// - we know the selected value (see selectedValue below)
+				// - we can use the attribute Type to tell us if they are the same
+				// - assuming entity is a SelfDirtinessTracker we can also know if the attribute is currently
+				//   considered dirty, and if really not dirty we would do the un-marking
+				// - of course that would mean a new method on SelfDirtinessTracker to allow un-marking
 			}
 			LOG.trace( "Done initializing lazy properties" );
 			return finalResult;
@@ -1734,19 +1737,30 @@ 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] );
 	}
 
-	// Used by Hibernate Reactive
+	/**
+	 * Used by Hibernate Reactive
+	 * @deprecated
+	 */
+	@Deprecated(since = "7.2", forRemoval = true)
 	protected boolean initializeLazyProperty(
 			final String fieldName,
 			final Object entity,
@@ -1763,6 +1777,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