Applied changes based on review suggestions

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

@@ -740,12 +740,9 @@ public class StatefulPersistenceContext implements PersistenceContext {
 			return narrowedProxy;
 		}
 		else {
-
 			if ( object != null ) {
 				final LazyInitializer li = ( (HibernateProxy) proxy ).getHibernateLazyInitializer();
-				if ( li.isUninitialized() ) {
+				assert !li.isUninitialized();
-					li.setImplementation( object );
-				}
 			}
 			return proxy;
 		}

hibernate-core/src/main/java/org/hibernate/loader/ast/internal/LoaderSelectBuilder.java

@@ -160,7 +160,7 @@ public class LoaderSelectBuilder {
 				sessionFactory,
 				loadable,
 				partsToSelect,
-				Arrays.asList( restrictedPart ),
+				Collections.singletonList( restrictedPart ),
 				cachedDomainResult,
 				numberOfKeysToLoad,
 				loadQueryInfluencers,

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/EmbeddableMappingType.java

@@ -528,21 +528,7 @@ public class EmbeddableMappingType implements ManagedMappingType, SelectionMappi
 		for ( int i = 0; i < attributes.size(); i++ ) {
 			final AttributeMapping attributeMapping = attributes.get( i );
 			final Object o = attributeMapping.getPropertyAccess().getGetter().get( value );
-			if ( attributeMapping instanceof ToOneAttributeMapping ) {
+			span += attributeMapping.forEachJdbcValue( o, clause, span + offset, consumer, session );
-				final ToOneAttributeMapping toOneAttributeMapping = (ToOneAttributeMapping) attributeMapping;
-				span += toOneAttributeMapping.getForeignKeyDescriptor().forEachJdbcValue(
-						toOneAttributeMapping.getAssociatedEntityMappingType()
-								.getIdentifierMapping()
-								.getIdentifier( o, session ),
-						clause,
-						span + offset,
-						consumer,
-						session
-				);
-			}
-			else {
-				span += attributeMapping.forEachJdbcValue( o, clause, span + offset, consumer, session );
-			}
 		}
 		return span;
 	}

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

@@ -9,6 +9,7 @@ package org.hibernate.metamodel.mapping.internal;
 import org.hibernate.LockMode;
 import org.hibernate.engine.FetchStrategy;
 import org.hibernate.engine.FetchTiming;
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.internal.util.StringHelper;
 import org.hibernate.mapping.ManyToOne;
 import org.hibernate.mapping.OneToOne;
@@ -27,6 +28,7 @@ import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.property.access.spi.PropertyAccess;
 import org.hibernate.query.EntityIdentifierNavigablePath;
 import org.hibernate.query.NavigablePath;
+import org.hibernate.sql.ast.Clause;
 import org.hibernate.sql.ast.SqlAstJoinType;
 import org.hibernate.sql.ast.spi.FromClauseAccess;
 import org.hibernate.sql.ast.spi.SqlAliasBase;
@@ -212,8 +214,22 @@ public class ToOneAttributeMapping extends AbstractSingularAttributeMapping
 		return navigableRole;
 	}
 
-	public boolean isForeignKeyOwner() {
+	@Override
-		return referencedPropertyName == null;
+	public int forEachJdbcValue(
+			Object value,
+			Clause clause,
+			int offset,
+			JdbcValuesConsumer consumer,
+			SharedSessionContractImplementor session) {
+		return getForeignKeyDescriptor().forEachJdbcValue(
+				getAssociatedEntityMappingType()
+						.getIdentifierMapping()
+						.getIdentifier( value, session ),
+				clause,
+				offset,
+				consumer,
+				session
+		);
 	}
 
 	@Override