PersistentOrderedSetIterator breaks if hash of an element changes #76

Author: Abduqodiri Qurbonzoda

core/commonMain/src/implementations/persistentOrderedMap/PersistentOrderedMapBuilderContentIterators.kt

@@ -28,7 +28,10 @@ internal open class PersistentOrderedMapBuilderLinksIterator<K, V>(
         lastIteratedKey = nextKey
         nextWasInvoked = true
         index++
-        val result = builder.hashMapBuilder[nextKey]!!
+        @Suppress("UNCHECKED_CAST")
+        val result = builder.hashMapBuilder.getOrElse(nextKey as K) {
+            throw ConcurrentModificationException("Hash code of a key ($nextKey) has changed after it was added to the persistent map.")
+        }
         nextKey = result.next
         return result
     }

core/commonMain/src/implementations/persistentOrderedMap/PersistentOrderedMapContentIterators.kt

@@ -21,7 +21,10 @@ internal open class PersistentOrderedMapLinksIterator<K, V>(
         if (!hasNext()) {
             throw NoSuchElementException()
         }
-        val result = hashMap[nextKey]!!
+        @Suppress("UNCHECKED_CAST")
+        val result = hashMap.getOrElse(nextKey as K) {
+            throw ConcurrentModificationException("Hash code of a key ($nextKey) has changed after it was added to the persistent map.")
+        }
         index++
         nextKey = result.next
         return result

core/commonMain/src/implementations/persistentOrderedSet/PersistentOrderedSetIterator.kt

@@ -19,7 +19,9 @@ internal open class PersistentOrderedSetIterator<E>(private var nextElement: Any
         @Suppress("UNCHECKED_CAST")
         val result = nextElement as E
         index++
-        nextElement = map[result]!!.next
+        nextElement = map.getOrElse(result) {
+            throw ConcurrentModificationException("Hash code of an element ($result) has changed after it was added to the persistent set.")
+        }.next
         return result
     }
 

core/commonTest/src/contract/map/ImmutableMapTest.kt

@@ -38,6 +38,24 @@ class ImmutableOrderedMapTest : ImmutableMapTest() {
         compare(listOf(1, 2), map.values) { collectionBehavior(ordered = true) }
         compare(mapOf("y" to 1, "x" to 2).entries, map.entries) { setBehavior(ordered = true) }
     }
+
+    @Test fun keyHashCodeChanged() {
+        val changing = mutableSetOf("ok")
+        val persistent: PersistentMap<Any, Any> = immutableMapOf("constant" to "fixed", changing to "modified")
+        assertEquals(1, persistent.keys.filter { it === changing }.size)
+        changing.add("break iteration")
+        assertFailsWith<ConcurrentModificationException> { persistent.keys.filter { it === changing } }
+    }
+
+    @Test fun builderKeyHashCodeChanged() {
+        val changing = mutableSetOf("ok")
+        val builder: PersistentMap.Builder<Any, Any> = immutableMapOf<Any, Any>().builder().apply {
+            putAll(listOf("constant" to "fixed", changing to "modified"))
+        }
+        assertEquals(1, builder.filter { it.key === changing }.size)
+        changing.add("break iteration")
+        assertFailsWith<ConcurrentModificationException> { builder.filter { it.key === changing } }
+    }
 }
 
 abstract class ImmutableMapTest {

core/commonTest/src/contract/set/ImmutableSetTest.kt

@@ -24,6 +24,24 @@ class ImmutableOrderedSetTest : ImmutableSetTestBase() {
     override fun <T> testBuilderToPersistentSet(builder: PersistentSet.Builder<T>) {
         assertSame(builder.build(), builder.toPersistentSet(), "toPersistent should call build()")
     }
+
+    @Test fun elementHashCodeChanged() {
+        val changing = mutableSetOf("ok")
+        val persistent: PersistentSet<Any> = immutableSetOf("constant", changing, "fix")
+        assertEquals(1, persistent.filter { it === changing }.size)
+        changing.add("break iteration")
+        assertFailsWith<ConcurrentModificationException> { persistent.filter { it === changing } }
+    }
+
+    @Test fun builderElementHashCodeChanged() {
+        val changing = mutableSetOf("ok")
+        val builder: PersistentSet.Builder<Any> = immutableSetOf<Any>().builder().apply {
+            addAll(listOf("constant", changing, "fix"))
+        }
+        assertEquals(1, builder.filter { it === changing }.size)
+        changing.add("break iteration")
+        assertFailsWith<ConcurrentModificationException> { builder.filter { it === changing } }
+    }
 }
 
 abstract class ImmutableSetTestBase {