Make PersistentHashSetBuilder iterator fail-fast

Author: Abduqodiri Qurbonzoda

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/immutableSet/PersistentHashSetBuilder.kt

@@ -23,7 +23,14 @@ internal class Marker
 internal class PersistentHashSetBuilder<E>(private var set: PersistentHashSet<E>) : AbstractMutableSet<E>(), PersistentSet.Builder<E> {
     internal var marker = Marker()
     internal var node = set.node
+    internal var modCount = 0
+
+    // Size change implies structural changes.
     override var size = set.size
+        set(value) {
+            field = value
+            modCount++
+        }
 
     override fun build(): PersistentSet<E> {
         set = if (node === set.node) {

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/immutableSet/PersistentHashSetMutableIterator.kt

@@ -18,31 +18,32 @@ package kotlinx.collections.immutable.implementations.immutableSet
 
 internal class PersistentHashSetMutableIterator<E>(private val builder: PersistentHashSetBuilder<E>)
     : PersistentHashSetIterator<E>(builder.node), MutableIterator<E> {
-    var lastReturned: E? = null
-    var nextWasInvoked = false
+    private var lastIteratedElement: E? = null
+    private var nextWasInvoked = false
+    private var expectedModCount = builder.modCount
 
     override fun next(): E {
+        checkForComodification()
         val next = super.next()
-        lastReturned = next
+        lastIteratedElement = next
         nextWasInvoked = true
         return next
     }
 
     override fun remove() {
-        if (!nextWasInvoked) {
-            throw NoSuchElementException()
-        }
+        checkNextWasInvoked()
         if (hasNext()) {
             val currentElement = currentElement()
 
-            assert(builder.remove(lastReturned))
+            assert(builder.remove(lastIteratedElement))
             resetPath(currentElement.hashCode(), builder.node, currentElement, 0)
         } else {
-            assert(builder.remove(lastReturned))
+            assert(builder.remove(lastIteratedElement))
         }
 
-        lastReturned = null
+        lastIteratedElement = null
         nextWasInvoked = false
+        expectedModCount = builder.modCount
     }
 
     private fun resetPath(hashCode: Int, node: TrieNode<*>, element: E, pathIndex: Int) {
@@ -63,12 +64,22 @@ internal class PersistentHashSetMutableIterator<E>(private val builder: Persiste
         if (cell is TrieNode<*>) {
             resetPath(hashCode, cell, element, pathIndex + 1)
         } else {
-            assert(cell == element)
+//            assert(cell == element)
             pathLastIndex = pathIndex
         }
     }
 
     private fun isCollision(node: TrieNode<*>): Boolean {
         return node.bitmap == 0
     }
+
+    private fun checkNextWasInvoked() {
+        if (!nextWasInvoked)
+            throw IllegalStateException()
+    }
+
+    private fun checkForComodification() {
+        if (builder.modCount != expectedModCount)
+            throw ConcurrentModificationException()
+    }
 }