Support null element and fail-fast iterator in OrderedSet and Builder

Author: Abduqodiri Qurbonzoda

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/persistentOrderedSet/PersistentOrderedSet.kt

@@ -18,13 +18,24 @@ package kotlinx.collections.immutable.implementations.persistentOrderedSet
 
 import kotlinx.collections.immutable.PersistentSet
 import kotlinx.collections.immutable.implementations.immutableMap.PersistentHashMap
+import kotlinx.collections.immutable.implementations.persistentOrderedMap.EndOfLink
 import kotlinx.collections.immutable.mutate
 
-internal class Links<out E>(val previous: E?, val next: E?)
+internal class Links(val previous: Any?, val next: Any?) {
+    constructor() : this(EndOfLink, EndOfLink)
+
+    fun withNext(newNext: Any?) = Links(previous, newNext)
+    fun withPrevious(newPrevious: Any?) = Links(newPrevious, next)
+
+    fun putNextLink(previous: Any?): Links {
+//        assert(next === EndOfLink)
+        return Links(previous, next)
+    }
+}
 
 internal class PersistentOrderedSet<E>(internal val firstElement: E?,
                                        internal val lastElement: E?,
-                                       internal val map: PersistentHashMap<E, Links<E>>): AbstractSet<E>(), PersistentSet<E> {
+                                       internal val map: PersistentHashMap<E, Links>): AbstractSet<E>(), PersistentSet<E> {
 
     override val size: Int
         get() = map.size
@@ -37,14 +48,16 @@ internal class PersistentOrderedSet<E>(internal val firstElement: E?,
         if (map.containsKey(element)) {
             return this
         }
-        if (lastElement == null) {  //  isEmpty
-            val newMap = map.put(element, Links<E>(null, null))
+        if (isEmpty()) {  //  isEmpty
+            val newMap = map.put(element, Links())
             return PersistentOrderedSet(element, element, newMap)
         }
-        val oldLinks = map[lastElement]!!
-        assert(oldLinks.next == null)
-        val newLinks = Links(oldLinks.previous, element)
-        val newMap = map.put(lastElement, newLinks).put(element, Links(lastElement, null))
+        val lastLinks = map[lastElement]!!
+//        assert(lastLinks.next === EndOfLink)
+
+        val newMap = map
+                .put(lastElement as E, lastLinks.withNext(element))
+                .put(element, lastLinks.putNextLink(lastElement))
         return PersistentOrderedSet(firstElement, element, newMap)
     }
 
@@ -56,18 +69,18 @@ internal class PersistentOrderedSet<E>(internal val firstElement: E?,
         val links = map[element] ?: return this
 
         var newMap = map.remove(element)
-        if (links.previous != null) {
+        if (links.previous !== EndOfLink) {
             val previousLinks = newMap[links.previous]!!
-            assert(previousLinks.next == element)
-            newMap = newMap.put(links.previous, Links(previousLinks.previous, links.next))
+//            assert(previousLinks.next == element)
+            newMap = newMap.put(links.previous as E, previousLinks.withNext(links.next))
         }
-        if (links.next != null) {
+        if (links.next !== EndOfLink) {
             val nextLinks = newMap[links.next]!!
-            assert(nextLinks.previous == element)
-            newMap = newMap.put(links.next, Links(links.previous, nextLinks.next))
+//            assert(nextLinks.previous == element)
+            newMap = newMap.put(links.next as E, nextLinks.withPrevious(links.previous))
         }
-        val newFirstElement = if (element == firstElement) links.next else firstElement
-        val newLastElement = if (element == lastElement) links.previous else lastElement
+        val newFirstElement = if (links.previous === EndOfLink) links.next as? E else firstElement
+        val newLastElement = if (links.next === EndOfLink) links.previous as? E else lastElement
         return PersistentOrderedSet(newFirstElement, newLastElement, newMap)
     }
 
@@ -92,7 +105,7 @@ internal class PersistentOrderedSet<E>(internal val firstElement: E?,
     }
 
     internal companion object {
-        private val EMPTY = PersistentOrderedSet<Nothing>(null, null, PersistentHashMap.emptyOf<Nothing, Links<Nothing>>())
+        private val EMPTY = PersistentOrderedSet(null, null, PersistentHashMap.emptyOf())
         internal fun <E> emptyOf(): PersistentSet<E> = EMPTY
     }
 }

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/persistentOrderedSet/PersistentOrderedSetBuilder.kt

@@ -17,6 +17,7 @@
 package kotlinx.collections.immutable.implementations.persistentOrderedSet
 
 import kotlinx.collections.immutable.PersistentSet
+import kotlinx.collections.immutable.implementations.persistentOrderedMap.EndOfLink
 
 internal class PersistentOrderedSetBuilder<E>(private var set: PersistentOrderedSet<E>) : AbstractMutableSet<E>(), PersistentSet.Builder<E> {
     internal var firstElement = set.firstElement
@@ -46,37 +47,38 @@ internal class PersistentOrderedSetBuilder<E>(private var set: PersistentOrdered
         if (mapBuilder.containsKey(element)) {
             return false
         }
-        if (lastElement == null) {  //  isEmpty
+        if (isEmpty()) {  //  isEmpty
             firstElement = element
             lastElement = element
-            mapBuilder[element] = Links<E>(null, null)
+            mapBuilder[element] = Links()
             return true
         }
-        val oldLinks = mapBuilder[lastElement!!]!!
-        assert(oldLinks.next == null)
-        val newLinks = Links(oldLinks.previous, element)
+        val lastLinks = mapBuilder[lastElement]!!
+//        assert(lastLinks.next === EndOfLink)
 
-        mapBuilder[lastElement!!] = newLinks
-        mapBuilder[element] = Links(lastElement, null)
+        mapBuilder[lastElement as E] = lastLinks.withNext(element)
+        mapBuilder[element] = lastLinks.putNextLink(lastElement)
         lastElement = element
         return true
     }
 
     override fun remove(element: E): Boolean {
         val links = mapBuilder.remove(element) ?: return false
 
-        if (links.previous != null) {
+        if (links.previous !== EndOfLink) {
             val previousLinks = mapBuilder[links.previous]!!
-            assert(previousLinks.next == element)
-            mapBuilder[links.previous] = Links(previousLinks.previous, links.next)
+//            assert(previousLinks.next == element)
+            mapBuilder[links.previous as E] = previousLinks.withNext(links.next)
+        } else {
+            firstElement = links.next as? E
         }
-        if (links.next != null) {
+        if (links.next !== EndOfLink) {
             val nextLinks = mapBuilder[links.next]!!
-            assert(nextLinks.previous == element)
-            mapBuilder[links.next] = Links(links.previous, nextLinks.next)
+//            assert(nextLinks.previous == element)
+            mapBuilder[links.next as E] = nextLinks.withPrevious(links.previous)
+        } else {
+            lastElement = links.previous as? E
         }
-        firstElement = if (element == firstElement) links.next else firstElement
-        lastElement = if (element == lastElement) links.previous else lastElement
         return true
     }
 

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/persistentOrderedSet/PersistentOrderedSetIterator.kt

@@ -17,20 +17,24 @@
 package kotlinx.collections.immutable.implementations.persistentOrderedSet
 
 internal open class PersistentOrderedSetIterator<E>(private var nextElement: E?,
-                                                    internal val map: Map<E, Links<E>>) : Iterator<E> {
+                                                    internal val map: Map<E, Links>) : Iterator<E> {
     internal var index = 0
 
     override fun hasNext(): Boolean {
         return index < map.size
     }
 
     override fun next(): E {
-        if (!hasNext())
-            throw NoSuchElementException()
+        checkHasNext()
 
         val result = nextElement as E
         index++
-        nextElement = map[result]!!.next
+        nextElement = map[result]!!.next as? E
         return result
     }
+
+    private fun checkHasNext() {
+        if (!hasNext())
+            throw NoSuchElementException()
+    }
 }

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/persistentOrderedSet/PersistentOrderedSetMutableIterator.kt

@@ -19,23 +19,34 @@ package kotlinx.collections.immutable.implementations.persistentOrderedSet
 internal class PersistentOrderedSetMutableIterator<E>(private val builder: PersistentOrderedSetBuilder<E>)
     : PersistentOrderedSetIterator<E>(builder.firstElement, builder.mapBuilder), MutableIterator<E> {
 
-    var lastReturned: E? = null
-    var nextWasInvoked = false
+    private var lastIteratedElement: E? = null
+    private var nextWasInvoked = false
+    private var expectedModCount = builder.mapBuilder.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()
-        }
-        builder.remove(lastReturned)
-        index--
-        lastReturned = null
+        checkNextWasInvoked()
+        builder.remove(lastIteratedElement)
+        lastIteratedElement = null
         nextWasInvoked = false
+        expectedModCount = builder.mapBuilder.modCount
+        index--
+    }
+
+    private fun checkNextWasInvoked() {
+        if (!nextWasInvoked)
+            throw IllegalStateException()
+    }
+
+    private fun checkForComodification() {
+        if (builder.mapBuilder.modCount != expectedModCount)
+            throw ConcurrentModificationException()
     }
 }