HHH-18271 Introduce lazy bitset for entities and other initializer improvements

* Cache more state in initializers
* Reduce mega-morphic call sites
* Do more efficient state resolving for query cache entries

Author: beikov

hibernate-core/src/main/java/org/hibernate/action/internal/EntityUpdateAction.java

@@ -266,6 +266,7 @@ private void handleGeneratedProperties(EntityEntry entry, GeneratedValues genera
 				nextVersion = getVersion( state, persister );
 			}
 			entry.postUpdate( instance, state, nextVersion );
+			entry.setMaybeLazySet( null );
 		}
 	}
 

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

@@ -27,10 +27,13 @@
 import org.hibernate.engine.spi.SessionImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.engine.spi.Status;
+import org.hibernate.internal.util.ImmutableBitSet;
 import org.hibernate.metamodel.mapping.AttributeMapping;
 import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.type.TypeHelper;
 
+import org.checkerframework.checker.nullness.qual.Nullable;
+
 import static org.hibernate.LockMode.PESSIMISTIC_FORCE_INCREMENT;
 import static org.hibernate.engine.internal.AbstractEntityEntry.BooleanState.EXISTS_IN_DATABASE;
 import static org.hibernate.engine.internal.AbstractEntityEntry.BooleanState.IS_BEING_REPLICATED;
@@ -71,6 +74,7 @@ public abstract class AbstractEntityEntry implements Serializable, EntityEntry {
 	protected transient EntityKey cachedEntityKey; // cached EntityKey (lazy-initialized)
 	protected final transient Object rowId;
 	protected final transient PersistenceContext persistenceContext;
+	protected transient @Nullable ImmutableBitSet maybeLazySet;
 	protected EntityEntryExtraState next;
 
 	/**
@@ -459,6 +463,16 @@ public void setReadOnly(boolean readOnly, Object entity) {
 		}
 	}
 
+	@Override
+	public @Nullable ImmutableBitSet getMaybeLazySet() {
+		return maybeLazySet;
+	}
+
+	@Override
+	public void setMaybeLazySet(@Nullable ImmutableBitSet maybeLazySet) {
+		this.maybeLazySet = maybeLazySet;
+	}
+
 	@Override
 	public String toString() {
 		return "EntityEntry"

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

@@ -9,10 +9,14 @@
 import java.io.IOException;
 import java.io.ObjectOutputStream;
 
+import org.hibernate.Internal;
 import org.hibernate.LockMode;
 import org.hibernate.collection.spi.PersistentCollection;
+import org.hibernate.internal.util.ImmutableBitSet;
 import org.hibernate.persister.entity.EntityPersister;
 
+import org.checkerframework.checker.nullness.qual.Nullable;
+
 /**
  * Information about the current state of a managed entity instance with respect
  * to its persistent state.
@@ -130,6 +134,16 @@ public interface EntityEntry {
 
 	void setReadOnly(boolean readOnly, Object entity);
 
+	/**
+	 * Has a bit set for every attribute position that is potentially lazy.
+	 * When {@code null}, no knowledge is available and every attribute must be assumed potentially lazy.
+	 */
+	@Internal
+	@Nullable ImmutableBitSet getMaybeLazySet();
+
+	@Internal
+	void setMaybeLazySet(@Nullable ImmutableBitSet maybeLazySet);
+
 	@Override
 	String toString();
 

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

@@ -0,0 +1,77 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html
+ */
+
+package org.hibernate.internal.util;
+
+import java.util.Arrays;
+import java.util.BitSet;
+
+/**
+ * An immutable variant of the {@link BitSet} class with some additional operations.
+ */
+public class ImmutableBitSet {
+
+	public static final ImmutableBitSet EMPTY = new ImmutableBitSet( new long[0] );
+
+	private static final int ADDRESS_BITS_PER_WORD = 6;
+
+	private final long[] words;
+
+	private static int wordIndex(int bitIndex) {
+		return bitIndex >> ADDRESS_BITS_PER_WORD;
+	}
+
+	private ImmutableBitSet(long[] words) {
+		this.words = words;
+	}
+
+	public static ImmutableBitSet valueOf(BitSet bitSet) {
+		final long[] words = bitSet.toLongArray();
+		return words.length == 0 ? EMPTY : new ImmutableBitSet( words );
+	}
+
+	public boolean get(int bitIndex) {
+		int wordIndex = wordIndex( bitIndex );
+		return wordIndex < words.length && ( ( words[wordIndex] & ( 1L << bitIndex ) ) != 0 );
+	}
+
+	public boolean isEmpty() {
+		return this == EMPTY;
+	}
+
+	public boolean contains(ImmutableBitSet set) {
+		if ( words.length < set.words.length ) {
+			return false;
+		}
+		// We check if every word from the given set is also contained in this set
+		for ( int i = 0; i < set.words.length; i++ ) {
+			if ( words[i] != ( set.words[i] | words[i] ) ) {
+				return false;
+			}
+		}
+		return true;
+	}
+
+	@Override
+	public int hashCode() {
+		return Arrays.hashCode( words );
+	}
+
+	@Override
+	public boolean equals(Object obj) {
+		if ( !( obj instanceof ImmutableBitSet ) ) {
+			return false;
+		}
+		if ( this == obj ) {
+			return true;
+		}
+
+		final ImmutableBitSet set = (ImmutableBitSet) obj;
+		return Arrays.equals( words, set.words );
+	}
+
+}

hibernate-core/src/main/java/org/hibernate/query/sqm/sql/internal/SqmMapEntryResult.java

@@ -102,5 +102,11 @@ public <X> void forEachResultAssembler(BiConsumer<Initializer<?>, X> consumer, X
 			keyAssembler.forEachResultAssembler( consumer, arg );
 			valueAssembler.forEachResultAssembler( consumer, arg );
 		}
+
+		@Override
+		public void resolveState(RowProcessingState rowProcessingState) {
+			keyAssembler.resolveState( rowProcessingState );
+			valueAssembler.resolveState( rowProcessingState );
+		}
 	}
 }

hibernate-core/src/main/java/org/hibernate/sql/results/graph/Initializer.java

@@ -64,7 +64,10 @@ public interface Initializer<Data extends InitializerData> {
 	ModelPart getInitializedPart();
 
 	default Object getResolvedInstance(Data data) {
-		return data.getState() == State.RESOLVED || data.getState() == State.INITIALIZED ? data.getInstance() : null;
+		assert data.getState() != State.UNINITIALIZED
+				&& data.getState() != State.KEY_RESOLVED
+				&& ( data.getState() != State.MISSING || data.getInstance() == null );
+		return data.getInstance();
 	}
 	default Object getResolvedInstance(RowProcessingState rowProcessingState) {
 		return getResolvedInstance( getData( rowProcessingState ) );
@@ -126,6 +129,12 @@ default void resolveInstance(RowProcessingState rowProcessingState) {
 		resolveInstance( getData( rowProcessingState ) );
 	}
 
+	void resolveState(Data data);
+
+	default void resolveState(RowProcessingState rowProcessingState) {
+		resolveState( getData( rowProcessingState ) );
+	}
+
 	/**
 	 * Step 2.2 - Use the given instance as resolved instance for this initializer.
 	 * Initializers are supposed to recursively call this method for sub-initializers.
@@ -206,6 +215,16 @@ static boolean isPartOfKey(NavigablePath navigablePath, InitializerParent<?> par
 				|| ForeignKeyDescriptor.TARGET_PART_NAME.equals( navigablePath.getLocalName() );
 	}
 
+	/**
+	 * Indicates whether calling resolve is needed when the object for this initializer is initialized already.
+	 */
+	boolean isEager();
+
+	/**
+	 * Indicates whether this initializers has eager sub-initializers, that could initialize lazy state of existing objects.
+	 */
+	boolean hasEagerSubInitializers();
+
 	/**
 	 * Indicates if this is a result or fetch initializer.
 	 */