Bug in PersistentMap equals implementation

core/commonMain/src/implementations/immutableMap/TrieNode.kt

@@ -181,7 +181,6 @@ internal class TrieNode<K, V>(
 
     /** The given [newNode] must not be a part of any persistent map instance. */
     private fun updateNodeAtIndex(nodeIndex: Int, positionMask: Int, newNode: TrieNode<K, V>, owner: MutabilityOwnership? = null): TrieNode<K, V> {
-//        assert(buffer[nodeIndex] !== newNode)
         val newNodeBuffer = newNode.buffer
         if (newNodeBuffer.size == 2 && newNode.nodeMap == 0) {
             if (buffer.size == 1) {
@@ -764,19 +763,17 @@ internal class TrieNode<K, V>(
             } else {
                 targetNode.mutableRemove(keyHash, key, shift + LOG_MAX_BRANCHING_FACTOR, mutator)
             }
-            return mutableReplaceNode(targetNode, newNode, nodeIndex, keyPositionMask, mutator.ownership)
+            return mutableReplaceNode(newNode, nodeIndex, keyPositionMask, mutator.ownership)
         }
 
         // key is absent
         return this
     }
 
-    private fun mutableReplaceNode(targetNode: TrieNode<K, V>, newNode: TrieNode<K, V>?, nodeIndex: Int, positionMask: Int, owner: MutabilityOwnership) = when {
+    private fun mutableReplaceNode(newNode: TrieNode<K, V>?, nodeIndex: Int, positionMask: Int, owner: MutabilityOwnership) = when {
         newNode == null ->
             mutableRemoveNodeAtIndex(nodeIndex, positionMask, owner)
-        targetNode !== newNode ->
-            updateNodeAtIndex(nodeIndex, positionMask, newNode, owner)
-        else -> this
+        else -> updateNodeAtIndex(nodeIndex, positionMask, newNode, owner)
     }
 
     fun remove(keyHash: Int, key: K, value: @UnsafeVariance V, shift: Int): TrieNode<K, V>? {
@@ -826,7 +823,7 @@ internal class TrieNode<K, V>(
             } else {
                 targetNode.mutableRemove(keyHash, key, value, shift + LOG_MAX_BRANCHING_FACTOR, mutator)
             }
-            return mutableReplaceNode(targetNode, newNode, nodeIndex, keyPositionMask, mutator.ownership)
+            return mutableReplaceNode(newNode, nodeIndex, keyPositionMask, mutator.ownership)
         }
 
         // key is absent

core/commonTest/src/stress/ObjectWrapper.kt → core/commonTest/src/ObjectWrapper.kt

@@ -3,7 +3,7 @@
  * Use of this source code is governed by the Apache 2.0 License that can be found in the LICENSE.txt file.
  */
 
-package tests.stress
+package tests
 
 import kotlinx.collections.immutable.internal.assert
 import kotlin.js.JsName

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

@@ -11,8 +11,8 @@ import tests.contract.compare
 import tests.contract.mapBehavior
 import tests.contract.setBehavior
 import tests.remove
-import tests.stress.IntWrapper
-import tests.stress.ObjectWrapper
+import tests.IntWrapper
+import tests.ObjectWrapper
 import kotlin.test.*
 
 class ImmutableHashMapTest : ImmutableMapTest() {

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

@@ -7,7 +7,7 @@ package tests.contract.map
 
 import kotlinx.collections.immutable.implementations.immutableMap.PersistentHashMap
 import kotlinx.collections.immutable.persistentHashMapOf
-import tests.stress.IntWrapper
+import tests.IntWrapper
 import kotlin.collections.iterator
 import kotlin.test.Test
 import kotlin.test.assertEquals

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

@@ -7,6 +7,7 @@ package tests.contract.map
 
 import kotlinx.collections.immutable.implementations.immutableMap.PersistentHashMap
 import kotlinx.collections.immutable.persistentHashMapOf
+import tests.IntWrapper
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertTrue
@@ -32,4 +33,44 @@ class PersistentHashMapTest {
         assertEquals(map3, map4.toMap())
         assertEquals(map3, map4)
     }
+
+    @Test
+    fun `builder should correctly handle multiple element removals in case of full collision`() {
+        val a = IntWrapper(0, 0)
+        val b = IntWrapper(1, 0)
+        val c = IntWrapper(2, 0)
+
+        val original: PersistentHashMap<IntWrapper, String> =
+            persistentHashMapOf(a to "a", b to "b", c to "c") as PersistentHashMap<IntWrapper, String>
+
+        val onlyA: PersistentHashMap<IntWrapper, String> =
+            persistentHashMapOf(a to "a") as PersistentHashMap<IntWrapper, String>
+
+        val builder = original.builder()
+        builder.remove(b)
+        builder.remove(c)
+        val removedBC = builder.build()
+
+        assertEquals(onlyA, removedBC)
+    }
+
+    @Test
+    fun `builder should correctly handle multiple element removals in case of partial collision`() {
+        val a = IntWrapper(0, 0)
+        val b = IntWrapper(1, 0)
+        val c = IntWrapper(2, 0)
+        val d = IntWrapper(3, 11)
+
+        val original: PersistentHashMap<IntWrapper, String> =
+            persistentHashMapOf(a to "a", b to "b", c to "c", d to "d") as PersistentHashMap<IntWrapper, String>
+
+        val afterImmutableRemoving = original.remove(b).remove(c)
+
+        val builder = original.builder()
+        builder.remove(b)
+        builder.remove(c)
+        val afterMutableRemoving = builder.build()
+
+        assertEquals(afterImmutableRemoving, afterMutableRemoving)
+    }
 }

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

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2016-2025 JetBrains s.r.o.
+ * Use of this source code is governed by the Apache 2.0 License that can be found in the LICENSE.txt file.
+ */
+
+package tests.contract.map
+
+import kotlinx.collections.immutable.minus
+import kotlinx.collections.immutable.persistentMapOf
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class PersistentOrderedMapTest {
+
+    /**
+     * Test from issue: https://github.com/Kotlin/kotlinx.collections.immutable/issues/198
+     */
+    @Test
+    fun `when removing multiple keys with identical hashcodes the remaining key should be correctly promoted`() {
+        class ChosenHashCode(
+            private val hashCode: Int,
+            private val name: String,
+        ) {
+            override fun equals(other: Any?): Boolean {
+                return other is ChosenHashCode && (other.name == name)
+            }
+
+            override fun hashCode(): Int {
+                return hashCode
+            }
+
+            override fun toString(): String {
+                return name
+            }
+        }
+
+        val a = ChosenHashCode(123, "A")
+        val b = ChosenHashCode(123, "B")
+        val c = ChosenHashCode(123, "C")
+
+        val abc = persistentMapOf(
+            a to "x",
+            b to "y",
+            c to "z",
+        )
+
+        val minusAb = abc.minus(arrayOf(a, b))
+        val cOnly = persistentMapOf(c to "z")
+
+        assertEquals(minusAb.entries, cOnly.entries)
+        assertEquals(minusAb, cOnly)
+    }
+}

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

@@ -10,7 +10,7 @@ import tests.contract.compare
 import tests.contract.setBehavior
 import tests.isDigit
 import tests.isUpperCase
-import tests.stress.IntWrapper
+import tests.IntWrapper
 import kotlin.test.*
 
 class ImmutableHashSetTest : ImmutableSetTestBase() {

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

@@ -7,7 +7,7 @@ package tests.contract.set
 
 import kotlinx.collections.immutable.implementations.immutableSet.PersistentHashSet
 import kotlinx.collections.immutable.persistentHashSetOf
-import tests.stress.IntWrapper
+import tests.IntWrapper
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith

core/commonTest/src/implementations/map/HashMapTrieNodeTest.kt

@@ -9,7 +9,7 @@ import kotlinx.collections.immutable.implementations.immutableMap.LOG_MAX_BRANCH
 import kotlinx.collections.immutable.implementations.immutableMap.MAX_SHIFT
 import kotlinx.collections.immutable.implementations.immutableMap.PersistentHashMap
 import kotlinx.collections.immutable.implementations.immutableMap.TrieNode
-import tests.stress.IntWrapper
+import tests.IntWrapper
 import kotlin.test.*
 
 class HashMapTrieNodeTest {

core/commonTest/src/stress/WrapperGenerator.kt

@@ -5,6 +5,7 @@
 
 package tests.stress
 
+import tests.ObjectWrapper
 import kotlin.random.Random
 
 

core/commonTest/src/stress/map/PersistentHashMapBuilderTest.kt

@@ -11,7 +11,7 @@ import tests.NForAlgorithmComplexity
 import tests.distinctStringValues
 import tests.remove
 import tests.stress.ExecutionTimeMeasuringTest
-import tests.stress.IntWrapper
+import tests.IntWrapper
 import tests.stress.WrapperGenerator
 import kotlin.random.Random
 import kotlin.test.*

core/commonTest/src/stress/map/PersistentHashMapTest.kt

@@ -11,7 +11,7 @@ import tests.NForAlgorithmComplexity
 import tests.distinctStringValues
 import tests.remove
 import tests.stress.ExecutionTimeMeasuringTest
-import tests.stress.IntWrapper
+import tests.IntWrapper
 import tests.stress.WrapperGenerator
 import kotlin.random.Random
 import kotlin.test.*

core/commonTest/src/stress/set/PersistentHashSetBuilderTest.kt

@@ -9,7 +9,7 @@ import kotlinx.collections.immutable.persistentHashSetOf
 import tests.NForAlgorithmComplexity
 import tests.distinctStringValues
 import tests.stress.ExecutionTimeMeasuringTest
-import tests.stress.IntWrapper
+import tests.IntWrapper
 import tests.stress.WrapperGenerator
 import kotlin.random.Random
 import kotlin.test.Test

core/commonTest/src/stress/set/PersistentHashSetTest.kt

@@ -9,7 +9,7 @@ import kotlinx.collections.immutable.persistentHashSetOf
 import tests.NForAlgorithmComplexity
 import tests.distinctStringValues
 import tests.stress.ExecutionTimeMeasuringTest
-import tests.stress.IntWrapper
+import tests.IntWrapper
 import tests.stress.WrapperGenerator
 import kotlin.random.Random
 import kotlin.test.Test