HHH-13458 Update Hibernate's custom IdentityMap to better match its use

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

@@ -232,7 +232,6 @@ public void clear() {
 		}
 
 		for ( Entry<Object, EntityEntry> objectEntityEntryEntry : entityEntryContext.reentrantSafeEntityEntries() ) {
-			// todo : I dont think this need be reentrant safe
 			if ( objectEntityEntryEntry.getKey() instanceof PersistentAttributeInterceptable ) {
 				final PersistentAttributeInterceptor interceptor = ( (PersistentAttributeInterceptable) objectEntityEntryEntry.getKey() ).$$_hibernate_getInterceptor();
 				if ( interceptor instanceof LazyAttributeLoadingInterceptor ) {
@@ -241,9 +240,8 @@ public void clear() {
 			}
 		}
 
-		for ( Map.Entry<PersistentCollection, CollectionEntry> aCollectionEntryArray : IdentityMap.concurrentEntries( collectionEntries ) ) {
-			aCollectionEntryArray.getKey().unsetSession( getSession() );
-		}
+		final SharedSessionContractImplementor session = getSession();
+		IdentityMap.onEachKey( collectionEntries, k -> k.unsetSession( session ) );
 
 		arrayHolders.clear();
 		entitiesByKey.clear();

hibernate-core/src/main/java/org/hibernate/event/internal/AbstractFlushingEventListener.java

@@ -193,7 +193,7 @@ private void prepareCollectionFlushes(PersistenceContext persistenceContext) thr
 		LOG.debug( "Dirty checking collections" );
 
 		for ( Map.Entry<PersistentCollection,CollectionEntry> entry :
-				IdentityMap.concurrentEntries( (Map<PersistentCollection,CollectionEntry>) persistenceContext.getCollectionEntries() )) {
+				IdentityMap.concurrentEntries( (Map<PersistentCollection,CollectionEntry>) persistenceContext.getCollectionEntries() ) ) {
 			entry.getValue().preFlush( entry.getKey() );
 		}
 	}
@@ -269,7 +269,7 @@ private int flushCollections(final EventSource session, final PersistenceContext
 
 		ActionQueue actionQueue = session.getActionQueue();
 		for ( Map.Entry<PersistentCollection,CollectionEntry> me :
-			IdentityMap.concurrentEntries( (Map<PersistentCollection,CollectionEntry>) persistenceContext.getCollectionEntries() )) {
+			IdentityMap.concurrentEntries( (Map<PersistentCollection,CollectionEntry>) persistenceContext.getCollectionEntries() ) ) {
 			PersistentCollection coll = me.getKey();
 			CollectionEntry ce = me.getValue();
 

hibernate-core/src/main/java/org/hibernate/internal/util/collections/IdentityMap.java

@@ -13,16 +13,17 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Consumer;
 
 /**
  * A <tt>Map</tt> where keys are compared by object identity,
  * rather than <tt>equals()</tt>.
  */
 public final class IdentityMap<K,V> implements Map<K,V> {
-	private final Map<IdentityKey<K>,V> map;
+
+	private final LinkedHashMap<IdentityKey<K>,V> map;
 	@SuppressWarnings( {"unchecked"})
-	private transient Entry<IdentityKey<K>,V>[] entryArray = new Entry[0];
-	private transient boolean dirty;
+	private transient Map.Entry<IdentityKey<K>,V>[] entryArray = null;
 
 	/**
 	 * Return a new instance of this class, with iteration
@@ -32,7 +33,7 @@ public final class IdentityMap<K,V> implements Map<K,V> {
 	 * @return The map
 	 */
 	public static <K,V> IdentityMap<K,V> instantiateSequenced(int size) {
-		return new IdentityMap<K,V>( new LinkedHashMap<IdentityKey<K>,V>( size ) );
+		return new IdentityMap<K,V>( new LinkedHashMap<>( size << 1, 0.6f ) );
 	}
 
 	/**
@@ -40,9 +41,8 @@ public static <K,V> IdentityMap<K,V> instantiateSequenced(int size) {
 	 *
 	 * @param underlyingMap The delegate map.
 	 */
-	private IdentityMap(Map<IdentityKey<K>,V> underlyingMap) {
+	private IdentityMap(LinkedHashMap<IdentityKey<K>,V> underlyingMap) {
 		map = underlyingMap;
-		dirty = true;
 	}
 
 	/**
@@ -57,6 +57,11 @@ public static <K,V> Map.Entry<K,V>[] concurrentEntries(Map<K,V> map) {
 		return ( (IdentityMap<K,V>) map ).entryArray();
 	}
 
+	public static <K,V> void onEachKey(Map<K,V> map, Consumer<K> consumer) {
+		final IdentityMap<K, V> identityMap = (IdentityMap<K, V>) map;
+		identityMap.map.forEach( (kIdentityKey, v) -> consumer.accept( kIdentityKey.key ) );
+	}
+
 	public Iterator<K> keyIterator() {
 		return new KeyIterator<K>( map.keySet().iterator() );
 	}
@@ -79,26 +84,27 @@ public boolean containsKey(Object key) {
 
 	@Override
 	public boolean containsValue(Object val) {
-		return map.containsValue(val);
+		throw new UnsupportedOperationException( "Avoid this operation: does not perform well" );
+		//return map.containsValue( val );
 	}
 
 	@Override
 	@SuppressWarnings( {"unchecked"})
 	public V get(Object key) {
-		return map.get( new IdentityKey(key) );
+		return map.get( new IdentityKey( key ) );
 	}
 
 	@Override
 	public V put(K key, V value) {
-		dirty = true;
-		return map.put( new IdentityKey<K>(key), value );
+		this.entryArray = null;
+		return map.put( new IdentityKey<K>( key ), value );
 	}
 
 	@Override
 	@SuppressWarnings( {"unchecked"})
 	public V remove(Object key) {
-		dirty = true;
-		return map.remove( new IdentityKey(key) );
+		this.entryArray = null;
+		return map.remove( new IdentityKey( key ) );
 	}
 
 	@Override
@@ -110,7 +116,6 @@ public void putAll(Map<? extends K, ? extends V> otherMap) {
 
 	@Override
 	public void clear() {
-		dirty = true;
 		entryArray = null;
 		map.clear();
 	}
@@ -118,6 +123,7 @@ public void clear() {
 	@Override
 	public Set<K> keySet() {
 		// would need an IdentitySet for this!
+		// (and we just don't use this method so it's ok)
 		throw new UnsupportedOperationException();
 	}
 
@@ -130,22 +136,21 @@ public Collection<V> values() {
 	public Set<Entry<K,V>> entrySet() {
 		Set<Entry<K,V>> set = new HashSet<Entry<K,V>>( map.size() );
 		for ( Entry<IdentityKey<K>, V> entry : map.entrySet() ) {
-			set.add( new IdentityMapEntry<K,V>( entry.getKey().getRealKey(), entry.getValue() ) );
+			set.add( new IdentityMapEntry<K,V>( entry.getKey().key, entry.getValue() ) );
 		}
 		return set;
 	}
 
 	@SuppressWarnings( {"unchecked"})
 	public Map.Entry[] entryArray() {
-		if (dirty) {
+		if ( entryArray == null ) {
 			entryArray = new Map.Entry[ map.size() ];
-			Iterator itr = map.entrySet().iterator();
-			int i=0;
+			final Iterator<Entry<IdentityKey<K>, V>> itr = map.entrySet().iterator();
+			int i = 0;
 			while ( itr.hasNext() ) {
-				Map.Entry me = (Map.Entry) itr.next();
-				entryArray[i++] = new IdentityMapEntry( ( (IdentityKey) me.getKey() ).key, me.getValue() );
+				final Entry<IdentityKey<K>, V> me = itr.next();
+				entryArray[i++] = new IdentityMapEntry( me.getKey().key, me.getValue() );
 			}
-			dirty = false;
 		}
 		return entryArray;
 	}
@@ -155,7 +160,7 @@ public String toString() {
 		return map.toString();
 	}
 
-	static final class KeyIterator<K> implements Iterator<K> {
+	private static final class KeyIterator<K> implements Iterator<K> {
 		private final Iterator<IdentityKey<K>> identityKeyIterator;
 
 		private KeyIterator(Iterator<IdentityKey<K>> iterator) {
@@ -167,7 +172,7 @@ public boolean hasNext() {
 		}
 
 		public K next() {
-			return identityKeyIterator.next().getRealKey();
+			return identityKeyIterator.next().key;
 		}
 
 		public void remove() {
@@ -175,9 +180,11 @@ public void remove() {
 		}
 
 	}
-		public static final class IdentityMapEntry<K,V> implements java.util.Map.Entry<K,V> {
+
+	private static final class IdentityMapEntry<K,V> implements java.util.Map.Entry<K,V> {
+
 		private final K key;
-		private V value;
+		private final V value;
 
 		IdentityMapEntry(final K key, final V value) {
 			this.key=key;
@@ -193,21 +200,16 @@ public V getValue() {
 		}
 
 		public V setValue(final V value) {
-			V result = this.value;
-			this.value = value;
-			return result;
+			throw new UnsupportedOperationException();
 		}
 	}
 
 	/**
-	 * We need to base the identity on {@link System#identityHashCode(Object)} but
-	 * attempt to lazily initialize and cache this value: being a native invocation
-	 * it is an expensive value to retrieve.
+	 * We need to base the identity on {@link System#identityHashCode(Object)}
 	 */
-	public static final class IdentityKey<K> implements Serializable {
+	private static final class IdentityKey<K> implements Serializable {
 
 		private final K key;
-		private int hash;
 
 		IdentityKey(K key) {
 			this.key = key;
@@ -221,21 +223,7 @@ public boolean equals(Object other) {
 
 		@Override
 		public int hashCode() {
-			if ( this.hash == 0 ) {
-				//We consider "zero" as non-initialized value
-				final int newHash = System.identityHashCode( key );
-				if ( newHash == 0 ) {
-					//So make sure we don't store zeros as it would trigger initialization again:
-					//any value is fine as long as we're deterministic.
-					this.hash = -1;
-					return -1;
-				}
-				else {
-					this.hash = newHash;
-					return newHash;
-				}
-			}
-			return hash;
+			return System.identityHashCode( key );
 		}
 
 		@Override
@@ -243,9 +231,6 @@ public String toString() {
 			return key.toString();
 		}
 
-		public K getRealKey() {
-			return key;
-		}
 	}
 
 }

hibernate-core/src/test/java/org/hibernate/test/customstructures/IdentityMapTest.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.test.customstructures;
+
+import java.util.Iterator;
+import java.util.Objects;
+
+import org.hibernate.internal.util.collections.IdentityMap;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class IdentityMapTest {
+
+    @Test
+    public void basicIdentityMapFunctionality() {
+        final IdentityMap<Holder, Object> map = IdentityMap.instantiateSequenced( 10 );
+        Holder k1 = new Holder( "k", 1 );
+        Holder s2 = new Holder( "s", 2 );
+        map.put( k1, "k1"  );
+        map.put( s2, "s2" );
+        map.put( k1, "K1!" );
+        Assert.assertEquals( 2, map.size() );
+        k1.name = "p";
+        Assert.assertEquals( 2, map.size() );
+        Assert.assertEquals( "K1!", map.get( k1 ) );
+        Holder k1similar = new Holder( "p", 1 );
+        map.put( k1similar, "notk1" );
+        Assert.assertEquals( "K1!", map.get( k1 ) );
+
+        IdentityMap.onEachKey( map, k -> k.value = 10 );
+
+        final Iterator<Holder> keyIterator = map.keyIterator();
+        int count = 0;
+        while ( keyIterator.hasNext() ) {
+            final Holder key = keyIterator.next();
+            Assert.assertNotNull( key );
+            count++;
+            Assert.assertEquals( 10, key.value );
+        }
+        Assert.assertEquals( 3, count );
+    }
+
+    private static class Holder {
+
+        //Evil: mutable keys!
+        private String name;
+        private int value;
+
+        public Holder(String name, int value) {
+            this.name = name;
+            this.value = value;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if ( this == o ) {
+                return true;
+            }
+            if ( o == null || getClass() != o.getClass() ) {
+                return false;
+            }
+            Holder holder = (Holder) o;
+            return value == holder.value &&
+                    name.equals( holder.name );
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash( name, value );
+        }
+    }
+}