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

@@ -596,7 +596,7 @@ private Fetch createDelayedCollectionFetch(
 		// returning a domain result assembler that returns LazyPropertyInitializer.UNFETCHED_PROPERTY
 		final var 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

@@ -71,6 +71,7 @@
 import org.hibernate.id.IdentifierGenerator;
 import org.hibernate.id.OptimizableGenerator;
 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;
@@ -1733,6 +1734,11 @@ protected boolean initializeLazyProperty(
 			final Object propValue) {
 		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
@@ -1767,6 +1773,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 var loadedState = entry.getLoadedState();
 		if ( loadedState != null ) {
 			// object has been loaded with setReadOnly(true); HHH-2236