Make PersistentHashMapBuilder iterators fail-fast

Author: Abduqodiri Qurbonzoda

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/immutableMap/PersistentHashMapBuilder.kt

@@ -24,7 +24,14 @@ internal class PersistentHashMapBuilder<K, V>(private var map: PersistentHashMap
     internal var marker = Marker()
     internal var node = map.node
     internal var operationResult: V? = null
+    internal var modCount = 0
+
+    // Size change implies structural changes.
     override var size = map.size
+        set(value) {
+            field = value
+            modCount++
+        }
 
     override fun build(): PersistentHashMap<K, V> {
         map = if (node === map.node) {

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/immutableMap/PersistentHashMapBuilderContentIterators.kt

@@ -43,30 +43,31 @@ internal abstract class PersistentHashMapBuilderBaseIterator<K, V, T>(private va
                                                                       path: Array<TrieNodeBaseIterator<K, V, T>>)
     : MutableIterator<T>, PersistentHashMapBaseIterator<K, V, T>(builder.node, path) {
 
-    private var lastKey: K? = null
+    private var lastIteratedKey: K? = null
     private var nextWasInvoked = false
+    private var expectedModCount = builder.modCount
 
     override fun next(): T {
-        lastKey = currentKey()
+        checkForComodification()
+        lastIteratedKey = currentKey()
         nextWasInvoked = true
         return super.next()
     }
 
     override fun remove() {
-        if (!nextWasInvoked) {
-            throw NoSuchElementException()
-        }
+        checkNextWasInvoked()
         if (hasNext()) {
             val currentKey = currentKey()
 
-            builder.remove(lastKey)
+            builder.remove(lastIteratedKey)
             resetPath(currentKey.hashCode(), builder.node, currentKey, 0)
         } else {
-            builder.remove(lastKey)
+            builder.remove(lastIteratedKey)
         }
 
-        lastKey = null
+        lastIteratedKey = null
         nextWasInvoked = false
+        expectedModCount = builder.modCount
     }
 
     private fun resetPath(keyHash: Int, node: TrieNode<*, *>, key: K, pathIndex: Int) {
@@ -76,13 +77,13 @@ internal abstract class PersistentHashMapBuilderBaseIterator<K, V, T>(private va
         if (node.hasDataAt(keyPosition)) { // key is directly in buffer
             val keyIndex = node.keyDataIndex(keyPosition)
 
-            assert(node.keyAtIndex(keyIndex) == key)
+//            assert(node.keyAtIndex(keyIndex) == key)
 
             path[pathIndex].reset(node.buffer, 2 * Integer.bitCount(node.dataMap), keyIndex)
             return
         }
 
-        assert(node.hasNodeAt(keyPosition)) // key is in node
+//        assert(node.hasNodeAt(keyPosition)) // key is in node
 
         val nodeIndex = node.keyNodeIndex(keyPosition)
         val targetNode = node.nodeAtIndex(nodeIndex)
@@ -96,6 +97,16 @@ internal abstract class PersistentHashMapBuilderBaseIterator<K, V, T>(private va
             resetPath(keyHash, targetNode, key, pathIndex + 1)
         }
     }
+
+    private fun checkNextWasInvoked() {
+        if (!nextWasInvoked)
+            throw IllegalStateException()
+    }
+
+    private fun checkForComodification() {
+        if (builder.modCount != expectedModCount)
+            throw ConcurrentModificationException()
+    }
 }
 
 internal class PersistentHashMapBuilderEntriesIterator<K, V>(builder: PersistentHashMapBuilder<K, V>)

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/immutableMap/PersistentHashMapContentIterators.kt

@@ -134,6 +134,7 @@ internal abstract class PersistentHashMapBaseIterator<K, V, T>(node: TrieNode<K,
     }
 
     protected fun currentKey(): K {
+        checkHasNext()
         return path[pathLastIndex].currentKey()
     }
 
@@ -142,11 +143,16 @@ internal abstract class PersistentHashMapBaseIterator<K, V, T>(node: TrieNode<K,
     }
 
     override fun next(): T {
-        assert(hasNext())
+        checkHasNext()
         val result = path[pathLastIndex].next()
         ensureNextEntryIsReady()
         return result
     }
+
+    private fun checkHasNext() {
+        if (!hasNext())
+            throw NoSuchElementException()
+    }
 }
 
 internal class PersistentHashMapEntriesIterator<K, V>(node: TrieNode<K, V>)

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/immutableMap/TrieNode.kt

@@ -19,7 +19,7 @@ internal class TrieNode<K, V>(var dataMap: Int,
         return dataMap and position != 0
     }
 
-    internal fun hasNodeAt(position: Int): Boolean {
+    private fun hasNodeAt(position: Int): Boolean {
         return nodeMap and position != 0
     }
 
@@ -31,7 +31,7 @@ internal class TrieNode<K, V>(var dataMap: Int,
         return buffer.size - 1 - Integer.bitCount(nodeMap and (position - 1))
     }
 
-    internal fun keyAtIndex(keyIndex: Int): K {
+    private fun keyAtIndex(keyIndex: Int): K {
         return buffer[keyIndex] as K
     }
 
@@ -84,11 +84,14 @@ internal class TrieNode<K, V>(var dataMap: Int,
     private fun mutableUpdateValueAtIndex(keyIndex: Int, value: V, mutator: PersistentHashMapBuilder<K, V>): TrieNode<K, V> {
 //        assert(buffer[keyIndex + 1] !== value)
 
-        mutator.operationResult = buffer[keyIndex + 1] as V
+        // If the [mutator] is exclusive owner of this node, update value at specified index in-place.
         if (marker === mutator.marker) {
             buffer[keyIndex + 1] = value
             return this
         }
+        // Structural change due to node replacement.
+        mutator.modCount++
+        // Create new node with updated value at specified index.
         val newBuffer = buffer.copyOf()
         newBuffer[keyIndex + 1] = value
         return TrieNode(dataMap, nodeMap, newBuffer, mutator.marker)
@@ -103,6 +106,8 @@ internal class TrieNode<K, V>(var dataMap: Int,
     }
 
     private fun mutableUpdateNodeAtIndex(nodeIndex: Int, newNode: TrieNode<K, V>, mutatorMarker: Marker): TrieNode<K, V> {
+//        assert(buffer[nodeIndex] !== newNode)
+
         if (marker === mutatorMarker) {
             buffer[nodeIndex] = newNode
             return this
@@ -286,19 +291,26 @@ internal class TrieNode<K, V>(var dataMap: Int,
     }
 
     private fun mutableCollisionPut(key: K, value: V, mutator: PersistentHashMapBuilder<K, V>): TrieNode<K, V> {
+        // Check if there is an entry with the specified key.
         for (i in 0 until buffer.size step ENTRY_SIZE) {
-            if (key == buffer[i]) {
+            if (key == buffer[i]) { // found entry with the specified key
                 mutator.operationResult = buffer[i + 1] as V
 
+                // If the [mutator] is exclusive owner of this node, update value of the entry in-place.
                 if (marker === mutator.marker) {
                     buffer[i + 1] = value
                     return this
                 }
+
+                // Structural change due to node replacement.
+                mutator.modCount++
+                // Create new node with updated entry value.
                 val newBuffer = buffer.copyOf()
                 newBuffer[i + 1] = value
                 return TrieNode(0, 0, newBuffer, mutator.marker)
             }
         }
+        // Create new collision node with the specified entry added to it.
         mutator.size++
         val newBuffer = bufferPutDataAtIndex(0, key, value)
         return TrieNode(0, 0, newBuffer, mutator.marker)
@@ -388,9 +400,9 @@ internal class TrieNode<K, V>(var dataMap: Int,
             val keyIndex = keyDataIndex(keyPosition)
 
             if (key == keyAtIndex(keyIndex)) {
+                modification.value = UPDATE_VALUE
                 if (valueAtKeyIndex(keyIndex) === value) { return this }
 
-                modification.value = UPDATE_VALUE
                 return updateValueAtIndex(keyIndex, value)
             }
             modification.value = PUT_KEY_VALUE
@@ -421,10 +433,9 @@ internal class TrieNode<K, V>(var dataMap: Int,
             val keyIndex = keyDataIndex(keyPosition)
 
             if (key == keyAtIndex(keyIndex)) {
-                if (valueAtKeyIndex(keyIndex) === value) { // needs discussion
-                    mutator.operationResult = value
-                    return this
-                }
+                mutator.operationResult = valueAtKeyIndex(keyIndex)
+                if (valueAtKeyIndex(keyIndex) === value) { return this }
+
                 return mutableUpdateValueAtIndex(keyIndex, value, mutator)
             }
             mutator.size++