Fix setValue method of MutableMapEntry

Author: Abduqodiri Qurbonzoda

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

@@ -8,32 +8,37 @@ package kotlinx.collections.immutable.implementations.immutableMap
 import kotlinx.collections.immutable.internal.assert
 
 
-internal class TrieNodeMutableEntriesIterator<K, V>(private val builder: PersistentHashMapBuilder<K, V>)
-    : TrieNodeBaseIterator<K, V, MutableMap.MutableEntry<K, V>>() {
+internal class TrieNodeMutableEntriesIterator<K, V>(
+        private val parentIterator: PersistentHashMapBuilderEntriesIterator<K, V>
+) : TrieNodeBaseIterator<K, V, MutableMap.MutableEntry<K, V>>() {
 
     override fun next(): MutableMap.MutableEntry<K, V> {
         assert(hasNextKey())
         index += 2
         @Suppress("UNCHECKED_CAST")
-        return MutableMapEntry(builder, buffer[index - 2] as K, buffer[index - 1] as V)
+        return MutableMapEntry(parentIterator, buffer[index - 2] as K, buffer[index - 1] as V)
     }
 }
 
-private class MutableMapEntry<K, V>(private val builder: PersistentHashMapBuilder<K, V>,
-                                    key: K,
-                                    override var value: V) : MapEntry<K, V>(key, value), MutableMap.MutableEntry<K, V> {
+private class MutableMapEntry<K, V>(
+        private val parentIterator: PersistentHashMapBuilderEntriesIterator<K, V>,
+        key: K,
+        override var value: V
+) : MapEntry<K, V>(key, value), MutableMap.MutableEntry<K, V> {
+
     override fun setValue(newValue: V): V {
         val result = value
         value = newValue
-        builder[key] = newValue
+        parentIterator.setValue(key, newValue)
         return result
     }
 }
 
 
-internal abstract class PersistentHashMapBuilderBaseIterator<K, V, T>(private val builder: PersistentHashMapBuilder<K, V>,
-                                                                      path: Array<TrieNodeBaseIterator<K, V, T>>)
-    : MutableIterator<T>, PersistentHashMapBaseIterator<K, V, T>(builder.node, path) {
+internal open class PersistentHashMapBuilderBaseIterator<K, V, T>(
+        private val builder: PersistentHashMapBuilder<K, V>,
+        path: Array<TrieNodeBaseIterator<K, V, T>>
+) : MutableIterator<T>, PersistentHashMapBaseIterator<K, V, T>(builder.node, path) {
 
     private var lastIteratedKey: K? = null
     private var nextWasInvoked = false
@@ -62,8 +67,33 @@ internal abstract class PersistentHashMapBuilderBaseIterator<K, V, T>(private va
         expectedModCount = builder.modCount
     }
 
+    fun setValue(key: K, newValue: V) {
+        if (!builder.containsKey(key)) return
+
+        if (hasNext()) {
+            val currentKey = currentKey()
+
+            builder[key] = newValue
+            resetPath(currentKey.hashCode(), builder.node, currentKey, 0)
+        } else {
+            builder[key] = newValue
+        }
+
+        expectedModCount = builder.modCount
+    }
+
     private fun resetPath(keyHash: Int, node: TrieNode<*, *>, key: K, pathIndex: Int) {
         val shift = pathIndex * LOG_MAX_BRANCHING_FACTOR
+
+        if (shift > MAX_SHIFT) {    // collision
+            path[pathIndex].reset(node.buffer, node.buffer.size, 0)
+            while (path[pathIndex].currentKey() != key) {
+                path[pathIndex].moveToNextKey()
+            }
+            pathLastIndex = pathIndex
+            return
+        }
+
         val keyPositionMask = 1 shl indexSegment(keyHash, shift)
 
         if (node.hasEntryAt(keyPositionMask)) { // key is directly in buffer
@@ -72,22 +102,16 @@ internal abstract class PersistentHashMapBuilderBaseIterator<K, V, T>(private va
 //            assert(node.keyAtIndex(keyIndex) == key)
 
             path[pathIndex].reset(node.buffer, ENTRY_SIZE * node.entryCount(), keyIndex)
+            pathLastIndex = pathIndex
             return
         }
 
-//        assert(node.hasNodeAt(keyPosition)) // key is in node
+//        assert(node.hasNodeAt(keyPositionMask)) // key is in node
 
         val nodeIndex = node.nodeIndex(keyPositionMask)
         val targetNode = node.nodeAtIndex(nodeIndex)
-        if (shift == MAX_SHIFT) {   // collision
-            path[pathIndex].reset(node.buffer, node.buffer.size, 0)
-            while (path[pathIndex].currentKey() != key) {
-                path[pathIndex].moveToNextKey()
-            }
-        } else {
-            path[pathIndex].reset(node.buffer, ENTRY_SIZE * node.entryCount(), nodeIndex)
-            resetPath(keyHash, targetNode, key, pathIndex + 1)
-        }
+        path[pathIndex].reset(node.buffer, ENTRY_SIZE * node.entryCount(), nodeIndex)
+        resetPath(keyHash, targetNode, key, pathIndex + 1)
     }
 
     private fun checkNextWasInvoked() {
@@ -101,8 +125,20 @@ internal abstract class PersistentHashMapBuilderBaseIterator<K, V, T>(private va
     }
 }
 
-internal class PersistentHashMapBuilderEntriesIterator<K, V>(builder: PersistentHashMapBuilder<K, V>)
-    : PersistentHashMapBuilderBaseIterator<K, V, MutableMap.MutableEntry<K, V>>(builder, Array(TRIE_MAX_HEIGHT + 1) { TrieNodeMutableEntriesIterator<K, V>(builder) })
+internal class PersistentHashMapBuilderEntriesIterator<K, V>(
+        builder: PersistentHashMapBuilder<K, V>
+) : MutableIterator<MutableMap.MutableEntry<K, V>> {
+    private val base = PersistentHashMapBuilderBaseIterator<K, V, MutableMap.MutableEntry<K, V>>(
+            builder,
+            Array(TRIE_MAX_HEIGHT + 1) { TrieNodeMutableEntriesIterator(this) }
+    )
+
+    override fun hasNext(): Boolean = base.hasNext()
+    override fun next(): MutableMap.MutableEntry<K, V> = base.next()
+    override fun remove(): Unit = base.remove()
+
+    fun setValue(key: K, newValue: V): Unit = base.setValue(key, newValue)
+}
 
 internal class PersistentHashMapBuilderKeysIterator<K, V>(builder: PersistentHashMapBuilder<K, V>)
     : PersistentHashMapBuilderBaseIterator<K, V, K>(builder, Array(TRIE_MAX_HEIGHT + 1) { TrieNodeKeysIterator<K, V>() })

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

@@ -103,7 +103,7 @@ internal abstract class PersistentHashMapBaseIterator<K, V, T>(
         protected val path: Array<TrieNodeBaseIterator<K, V, T>>
 ) : Iterator<T> {
 
-    private var pathLastIndex = 0
+    protected var pathLastIndex = 0
     @JsName("_hasNext")
     private var hasNext = true