HHH-19805 Update unfetched attributes of already initialized entity when data is fetched

Author: beikov

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