Replace ModificationWrapper with ModificationResult

It's created only when some change to the map is made.

Author: ilya-g

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

@@ -18,12 +18,6 @@ package kotlinx.collections.immutable.implementations.immutableMap
 
 import kotlinx.collections.immutable.*
 
-internal const val NO_MODIFICATION = 0
-internal const val UPDATE_VALUE = 1
-internal const val PUT_KEY_VALUE = 2
-internal class ModificationWrapper(var value: Int = NO_MODIFICATION)
-
-
 internal class PersistentHashMap<K, V>(internal val node: TrieNode<K, V>,
                                        override val size: Int): AbstractMap<K, V>(), PersistentMap<K, V> {
 
@@ -61,11 +55,8 @@ internal class PersistentHashMap<K, V>(internal val node: TrieNode<K, V>,
     }
 
     override fun put(key: K, value: @UnsafeVariance V): PersistentHashMap<K, V> {
-        val modification = ModificationWrapper()
-        val newNode = node.put(key.hashCode(), key, value, 0, modification)
-        if (node === newNode) { return this }
-        val sizeDelta = if (modification.value == PUT_KEY_VALUE) 1 else 0
-        return PersistentHashMap(newNode, size + sizeDelta)
+        val newNodeResult = node.put(key.hashCode(), key, value, 0) ?: return this
+        return PersistentHashMap(newNodeResult.node, size + newNodeResult.sizeDelta)
     }
 
     override fun remove(key: K): PersistentHashMap<K, V> {

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

@@ -75,6 +75,14 @@ internal class TrieNode<K, V>(
 ) {
     constructor(dataMap: Int, nodeMap: Int, buffer: Array<Any?>) : this(dataMap, nodeMap, buffer, null)
 
+    internal class ModificationResult<K, V>(var node: TrieNode<K, V>, val sizeDelta: Int) {
+        inline fun replaceNode(operation: (TrieNode<K, V>) -> TrieNode<K, V>): ModificationResult<K, V> =
+                apply { node = operation(node) }
+    }
+
+    private fun asInsertResult() = ModificationResult(this, 1)
+    private fun asUpdateResult() = ModificationResult(this, 0)
+
     internal var buffer: Array<Any?> = buffer
         private set
 
@@ -333,21 +341,19 @@ internal class TrieNode<K, V>(
         return null
     }
 
-    private fun collisionPut(key: K, value: V, modification: ModificationWrapper): TrieNode<K, V> {
+    private fun collisionPut(key: K, value: V): ModificationResult<K, V>? {
         for (i in 0 until buffer.size step ENTRY_SIZE) {
             if (key == buffer[i]) {
                 if (value === buffer[i + 1]) {
-                    return this
+                    return null
                 }
-                modification.value = UPDATE_VALUE
                 val newBuffer = buffer.copyOf()
                 newBuffer[i + 1] = value
-                return TrieNode(0, 0, newBuffer)
+                return TrieNode<K, V>(0, 0, newBuffer).asUpdateResult()
             }
         }
-        modification.value = PUT_KEY_VALUE
         val newBuffer = buffer.insertEntryAtIndex(0, key, value)
-        return TrieNode(0, 0, newBuffer)
+        return TrieNode<K, V>(0, 0, newBuffer).asInsertResult()
     }
 
     private fun mutableCollisionPut(key: K, value: V, mutator: PersistentHashMapBuilder<K, V>): TrieNode<K, V> {
@@ -454,37 +460,33 @@ internal class TrieNode<K, V>(
         return null
     }
 
-    fun put(keyHash: Int, key: K, value: @UnsafeVariance V, shift: Int, modification: ModificationWrapper): TrieNode<K, V> {
+    fun put(keyHash: Int, key: K, value: @UnsafeVariance V, shift: Int): ModificationResult<K, V>? {
         val keyPositionMask = 1 shl indexSegment(keyHash, shift)
 
         if (hasEntryAt(keyPositionMask)) { // key is directly in buffer
             val keyIndex = entryKeyIndex(keyPositionMask)
 
             if (key == keyAtIndex(keyIndex)) {
-                modification.value = UPDATE_VALUE
-                if (valueAtKeyIndex(keyIndex) === value) return this
+                if (valueAtKeyIndex(keyIndex) === value) return null
 
-                return updateValueAtIndex(keyIndex, value)
+                return updateValueAtIndex(keyIndex, value).asUpdateResult()
             }
-            modification.value = PUT_KEY_VALUE
-            return moveEntryToNode(keyIndex, keyPositionMask, keyHash, key, value, shift)
+            return moveEntryToNode(keyIndex, keyPositionMask, keyHash, key, value, shift).asInsertResult()
         }
         if (hasNodeAt(keyPositionMask)) { // key is in node
             val nodeIndex = nodeIndex(keyPositionMask)
 
             val targetNode = nodeAtIndex(nodeIndex)
-            val newNode = if (shift == MAX_SHIFT) {
-                targetNode.collisionPut(key, value, modification)
+            val putResult = if (shift == MAX_SHIFT) {
+                targetNode.collisionPut(key, value) ?: return null
             } else {
-                targetNode.put(keyHash, key, value, shift + LOG_MAX_BRANCHING_FACTOR, modification)
+                targetNode.put(keyHash, key, value, shift + LOG_MAX_BRANCHING_FACTOR) ?: return null
             }
-            if (targetNode === newNode) return this
-            return updateNodeAtIndex(nodeIndex, newNode)
+            return putResult.replaceNode { node -> updateNodeAtIndex(nodeIndex, node) }
         }
 
         // no entry at this key hash segment
-        modification.value = PUT_KEY_VALUE
-        return insertEntryAt(keyPositionMask, key, value)
+        return insertEntryAt(keyPositionMask, key, value).asInsertResult()
     }
 
     fun mutablePut(keyHash: Int, key: K, value: @UnsafeVariance V, shift: Int, mutator: PersistentHashMapBuilder<K, V>): TrieNode<K, V> {