Support null key and fail-fast iterators in OrderedMap and Builder

Author: Abduqodiri Qurbonzoda

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/persistentOrderedMap/PersistentOrderedMap.kt

@@ -22,11 +22,24 @@ import kotlinx.collections.immutable.PersistentMap
 import kotlinx.collections.immutable.implementations.immutableMap.PersistentHashMap
 import kotlinx.collections.immutable.mutate
 
-internal class LinkedValue<out K, out V>(val value: V, val previous: K?, val next: K?)
+internal object EndOfLink
+
+internal class LinkedValue<V>(val value: V, val previous: Any?, val next: Any?) {
+    constructor(value: V) : this(value, EndOfLink, EndOfLink)
+
+    fun withValue(newValue: V) = LinkedValue(newValue, previous, next)
+    fun withPrevious(newPrevious: Any?) = LinkedValue(value, newPrevious, next)
+    fun withNext(newNext: Any?) = LinkedValue(value, previous, newNext)
+
+    fun putNextLink(value: V, previous: Any?): LinkedValue<V> {
+//        assert(next === EndOfLink)
+        return LinkedValue(value, previous, EndOfLink)
+    }
+}
 
 internal class PersistentOrderedMap<K, V>(internal val firstKey: K?,
                                           internal val lastKey: K?,
-                                          internal val map: PersistentHashMap<K, LinkedValue<K, V>>): AbstractMap<K, V>(), PersistentMap<K, V> {
+                                          internal val map: PersistentHashMap<K, LinkedValue<V>>): AbstractMap<K, V>(), PersistentMap<K, V> {
 
     override val size: Int
         get() = map.size
@@ -70,36 +83,37 @@ internal class PersistentOrderedMap<K, V>(internal val firstKey: K?,
             if (links.value == value) {
                 return this
             }
-            val newMap = map.put(key, LinkedValue(value, links.previous, links.next))
+            val newMap = map.put(key, links.withValue(value))
             return PersistentOrderedMap(firstKey, lastKey, newMap)
         }
         if (isEmpty()) {
-            val newMap = map.put(key, LinkedValue<K, V>(value, null, null))
+            val newMap = map.put(key, LinkedValue(value))
             return PersistentOrderedMap(key, key, newMap)
         }
-        val oldLinks = map[lastKey]!!
-        assert(oldLinks.next == null)
-        val newLinks = LinkedValue(oldLinks.value, oldLinks.previous, key)
-        val newMap = map.put(lastKey as K, newLinks).put(key, LinkedValue(value, lastKey, null))
+        val lastLink = map[lastKey]!!
+//        assert(lastLink.next === EndOfLink)
+        val newMap = map
+                .put(lastKey as K, lastLink.withNext(key))
+                .put(key, lastLink.putNextLink(value, lastKey))
         return PersistentOrderedMap(firstKey, key, newMap)
     }
 
     override fun remove(key: K): PersistentOrderedMap<K, V> {
         val links = map[key] ?: return this
 
         var newMap = map.remove(key)
-        if (key != firstKey) {
+        if (links.previous !== EndOfLink) {
             val previousLinks = newMap[links.previous]!!
-            assert(previousLinks.next == key)
-            newMap = newMap.put(links.previous as K, LinkedValue(previousLinks.value, previousLinks.previous, links.next))
+//            assert(previousLinks.next == key)
+            newMap = newMap.put(links.previous as K, previousLinks.withNext(links.next))
         }
-        if (key != lastKey) {
+        if (links.next !== EndOfLink) {
             val nextLinks = newMap[links.next]!!
-            assert(nextLinks.previous == key)
-            newMap = newMap.put(links.next as K, LinkedValue(nextLinks.value, links.previous, nextLinks.next))
+//            assert(nextLinks.previous == key)
+            newMap = newMap.put(links.next as K, nextLinks.withPrevious(links.previous))
         }
-        val newFirstKey = if (key == firstKey) links.next else firstKey
-        val newLastKey = if (key == lastKey) links.previous else lastKey
+        val newFirstKey = if (links.previous === EndOfLink) links.next as? K else firstKey
+        val newLastKey = if (links.next === EndOfLink) links.previous as? K else lastKey
         return PersistentOrderedMap(newFirstKey, newLastKey, newMap)
     }
 

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/persistentOrderedMap/PersistentOrderedMapBuilder.kt

@@ -64,44 +64,45 @@ internal class PersistentOrderedMapBuilder<K, V>(private var map: PersistentOrde
     override fun put(key: K, value: @UnsafeVariance V): V? {
         val links = mapBuilder[key]
         if (links != null) {
-            if (links.value == value) {
+            if (links.value === value) {
                 return value
             }
-            mapBuilder[key] = LinkedValue(value, links.previous, links.next)
+            mapBuilder[key] = links.withValue(value)
             return links.value
         }
 
         if (isEmpty()) {  //  isEmpty
             firstKey = key
             lastKey = key
-            mapBuilder[key] = LinkedValue<K, V>(value, null, null)
+            mapBuilder[key] = LinkedValue(value)
             return null
         }
-        val oldLinks = mapBuilder[lastKey]!!
-        assert(oldLinks.next == null)
-        val newLinks = LinkedValue(oldLinks.value, oldLinks.previous, key)
+        val lastLink = mapBuilder[lastKey]!!
+        assert(lastLink.next == EndOfLink)
 
-        mapBuilder[lastKey as K] = newLinks
-        mapBuilder[key] = LinkedValue(value, lastKey, null)
+        mapBuilder[lastKey as K] = lastLink.withNext(key)
+        mapBuilder[key] = lastLink.putNextLink(value, lastKey)
         lastKey = key
         return null
     }
 
     override fun remove(key: K): V? {
         val links = mapBuilder.remove(key) ?: return null
 
-        if (links.previous != null) {
+        if (links.previous !== EndOfLink) {
             val previousLinks = mapBuilder[links.previous]!!
-            assert(previousLinks.next == key)
-            mapBuilder[links.previous] = LinkedValue(previousLinks.value, previousLinks.previous, links.next)
+//            assert(previousLinks.next == key)
+            mapBuilder[links.previous as K] = previousLinks.withNext(links.next)
+        } else {
+            firstKey = links.next as? K
         }
-        if (links.next != null) {
+        if (links.next !== EndOfLink) {
             val nextLinks = mapBuilder[links.next]!!
-            assert(nextLinks.previous == key)
-            mapBuilder[links.next] = LinkedValue(nextLinks.value, links.previous, nextLinks.next)
+//            assert(nextLinks.previous == key)
+            mapBuilder[links.next as K] = nextLinks.withPrevious(links.previous)
+        } else {
+            lastKey = links.previous as? K
         }
-        firstKey = if (key == firstKey) links.next else firstKey
-        lastKey = if (key == lastKey) links.previous else lastKey
         return links.value
     }
 

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/persistentOrderedMap/PersistentOrderedMapBuilderContentIterators.kt

@@ -18,80 +18,97 @@ package kotlinx.collections.immutable.implementations.persistentOrderedMap
 
 import kotlinx.collections.immutable.implementations.immutableMap.MapEntry
 
-internal open class PersistentOrderedMapBuilderLinksIterator<K, V>(internal var nextKey: K?,
-                                                                   internal val map: MutableMap<K, LinkedValue<K, V>>) : MutableIterator<LinkedValue<K, V>> {
-    internal var lastProcessedKey: K? = null
+internal open class PersistentOrderedMapBuilderLinksIterator<K, V>(
+        private var nextKey: K?,
+        internal val builder: PersistentOrderedMapBuilder<K, V>
+) : MutableIterator<LinkedValue<V>> {
+
+    internal var lastIteratedKey: K? = null
     private var nextWasInvoked = false
+    private var expectedModCount = builder.mapBuilder.modCount
     internal var index = 0
 
     override fun hasNext(): Boolean {
-        return index < map.size
+        return index < builder.size
     }
 
-    override fun next(): LinkedValue<K, V> {
-        if (!hasNext()) {
-            throw NoSuchElementException()
-        }
-        lastProcessedKey = nextKey
+    override fun next(): LinkedValue<V> {
+        checkForComodification()
+        checkHasNext()
+        lastIteratedKey = nextKey
         nextWasInvoked = true
         index++
-        val result = map[nextKey as K]!!
-        nextKey = result.next
+        val result = builder.mapBuilder[nextKey]!!
+        nextKey = result.next as? K
         return result
     }
 
     override fun remove() {
-        if (!nextWasInvoked) {
-            throw NoSuchElementException()
-        }
-        map.remove(lastProcessedKey)
-        lastProcessedKey = null
+        checkNextWasInvoked()
+        builder.remove(lastIteratedKey)
+        lastIteratedKey = null
         nextWasInvoked = false
+        expectedModCount = builder.mapBuilder.modCount
         index--
     }
+
+    private fun checkHasNext() {
+        if (!hasNext())
+            throw NoSuchElementException()
+    }
+
+    private fun checkNextWasInvoked() {
+        if (!nextWasInvoked)
+            throw IllegalStateException()
+    }
+
+    private fun checkForComodification() {
+        if (builder.mapBuilder.modCount != expectedModCount)
+            throw ConcurrentModificationException()
+    }
 }
 
 internal class PersistentOrderedMapBuilderEntriesIterator<K, V>(map: PersistentOrderedMapBuilder<K, V>): MutableIterator<MutableMap.MutableEntry<K, V>> {
-    private val internal = PersistentOrderedMapBuilderLinksIterator(map.firstKey, map.mapBuilder)
+    private val internal = PersistentOrderedMapBuilderLinksIterator(map.firstKey, map)
 
     override fun hasNext(): Boolean {
         return internal.hasNext()
     }
 
     override fun next(): MutableMap.MutableEntry<K, V> {
         val links = internal.next()
-        return MutableMapEntry(internal.map, internal.lastProcessedKey as K, links)
+        return MutableMapEntry(internal.builder.mapBuilder, internal.lastIteratedKey as K, links)
     }
 
     override fun remove() {
         internal.remove()
     }
 }
 
-private class MutableMapEntry<K, V>(private val mutableMap: MutableMap<K, LinkedValue<K, V>>,
+private class MutableMapEntry<K, V>(private val mutableMap: MutableMap<K, LinkedValue<V>>,
                                     key: K,
-                                    private var links: LinkedValue<K, V>) : MapEntry<K, V>(key, links.value), MutableMap.MutableEntry<K, V> {
+                                    private var links: LinkedValue<V>) : MapEntry<K, V>(key, links.value), MutableMap.MutableEntry<K, V> {
     override val value: V
         get() = links.value
 
     override fun setValue(newValue: V): V {
         val result = links.value
-        links = LinkedValue(newValue, links.previous, links.next)
+        links = links.withValue(newValue)
         mutableMap[key] = links
         return result
     }
 }
 
 internal class PersistentOrderedMapBuilderKeysIterator<out K, out V>(map: PersistentOrderedMapBuilder<K, V>): MutableIterator<K> {
-    private val internal = PersistentOrderedMapBuilderLinksIterator(map.firstKey, map.mapBuilder)
+    private val internal = PersistentOrderedMapBuilderLinksIterator(map.firstKey, map)
 
     override fun hasNext(): Boolean {
         return internal.hasNext()
     }
 
     override fun next(): K {
         internal.next()
-        return internal.lastProcessedKey as K
+        return internal.lastIteratedKey as K
     }
 
     override fun remove() {
@@ -100,7 +117,7 @@ internal class PersistentOrderedMapBuilderKeysIterator<out K, out V>(map: Persis
 }
 
 internal class PersistentOrderedMapBuilderValuesIterator<out K, out V>(map: PersistentOrderedMapBuilder<K, V>): MutableIterator<V> {
-    private val internal = PersistentOrderedMapBuilderLinksIterator(map.firstKey, map.mapBuilder)
+    private val internal = PersistentOrderedMapBuilderLinksIterator(map.firstKey, map)
 
     override fun hasNext(): Boolean {
         return internal.hasNext()

kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/persistentOrderedMap/PersistentOrderedMapContentIterators.kt

@@ -18,21 +18,21 @@ package kotlinx.collections.immutable.implementations.persistentOrderedMap
 
 import kotlinx.collections.immutable.implementations.immutableMap.MapEntry
 
-internal open class PersistentOrderedMapLinksIterator<K, out V>(internal var nextKey: K?,
-                                                                internal val map: Map<K, LinkedValue<K, V>>) : Iterator<LinkedValue<K, V>> {
+internal open class PersistentOrderedMapLinksIterator<K, V>(internal var nextKey: K?,
+                                                            internal val map: Map<K, LinkedValue<V>>) : Iterator<LinkedValue<V>> {
     internal var index = 0
 
     override fun hasNext(): Boolean {
         return index < map.size
     }
 
-    override fun next(): LinkedValue<K, V> {
+    override fun next(): LinkedValue<V> {
         if (!hasNext()) {
             throw NoSuchElementException()
         }
-        val result = map[nextKey as K]!!
+        val result = map[nextKey]!!
         index++
-        nextKey = result.next
+        nextKey = result.next as? K
         return result
     }