Do note create node containing a single other node

Kotlin · Oct 9, 2019 · 3eb424c · 3eb424c
1 parent 5dde861
commit 3eb424c
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 6 deletions.
diff --git a/core/commonMain/src/implementations/immutableMap/TrieNode.kt b/core/commonMain/src/implementations/immutableMap/TrieNode.kt
@@ -172,10 +172,17 @@ internal class TrieNode<K, V>(
         return TrieNode(dataMap, nodeMap, newBuffer, mutator.ownership)
     }
 
+    /** The given [newNode] must not be a part of any persistent map instance. */
     private fun updateNodeAtIndex(nodeIndex: Int, positionMask: Int, newNode: TrieNode<K, V>): TrieNode<K, V> {
 //        assert(buffer[nodeIndex] !== newNode)
         val newNodeBuffer = newNode.buffer
         if (newNodeBuffer.size == 2 && newNode.nodeMap == 0) {
+            if (buffer.size == 1) {
+//                assert(dataMap == 0 && nodeMap xor positionMask == 0)
+                newNode.dataMap = nodeMap
+                return newNode
+            }
+
             val keyIndex = entryKeyIndex(positionMask)
             val newBuffer = buffer.replaceNodeWithEntry(nodeIndex, keyIndex, newNodeBuffer[0], newNodeBuffer[1])
             return TrieNode(dataMap xor positionMask, nodeMap xor positionMask, newBuffer)
@@ -186,11 +193,18 @@ internal class TrieNode<K, V>(
         return TrieNode(dataMap, nodeMap, newBuffer)
     }
 
+    /** The given [newNode] must not be a part of any persistent map instance. */
     private fun mutableUpdateNodeAtIndex(nodeIndex: Int, positionMask: Int, newNode: TrieNode<K, V>, owner: MutabilityOwnership): TrieNode<K, V> {
 //        assert(buffer[nodeIndex] !== newNode)
 
         val newNodeBuffer = newNode.buffer
         if (newNodeBuffer.size == 2 && newNode.nodeMap == 0) {
+            if (buffer.size == 1) {
+//                assert(dataMap == 0 && nodeMap xor positionMask == 0)
+                newNode.dataMap = nodeMap
+                return newNode
+            }
+
             val keyIndex = entryKeyIndex(positionMask)
             val newBuffer = buffer.replaceNodeWithEntry(nodeIndex, keyIndex, newNodeBuffer[0], newNodeBuffer[1])
 

diff --git a/core/commonMain/src/implementations/immutableSet/TrieNode.kt b/core/commonMain/src/implementations/immutableSet/TrieNode.kt
@@ -89,29 +89,39 @@ internal class TrieNode<E>(
         return TrieNode(bitmap or positionMask, newBuffer, owner)
     }
 
+    /** The given [newNode] must not be a part of any persistent set instance. */
     private fun updateNodeAtIndex(nodeIndex: Int, newNode: TrieNode<E>): TrieNode<E> {
 //        assert(buffer[nodeIndex] !== newNode)
+        val cell: Any?
 
-        val newBuffer = buffer.copyOf()
-
-        // TODO: check how this changes affect `add` operation performance.
-        // Try to not create this newNode, but pass here the remained element.
         val newNodeBuffer = newNode.buffer
         if (newNodeBuffer.size == 1 && newNodeBuffer[0] !is TrieNode<*>) {
-            newBuffer[nodeIndex] = newNodeBuffer[0]
+            if (buffer.size == 1) {
+                newNode.bitmap = bitmap
+                return newNode
+            }
+            cell = newNodeBuffer[0]
         } else {
-            newBuffer[nodeIndex] = newNode
+            cell = newNode
         }
 
+        val newBuffer = buffer.copyOf()
+        newBuffer[nodeIndex] = cell
         return TrieNode(bitmap, newBuffer)
     }
 
+    /** The given [newNode] must not be a part of any persistent set instance. */
     private fun mutableUpdateNodeAtIndex(nodeIndex: Int, newNode: TrieNode<E>, owner: MutabilityOwnership): TrieNode<E> {
 //        assert(buffer[nodeIndex] !== newNode)
 
         val cell: Any?
+
         val newNodeBuffer = newNode.buffer
         if (newNodeBuffer.size == 1 && newNodeBuffer[0] !is TrieNode<*>) {
+            if (buffer.size == 1) {
+                newNode.bitmap = bitmap
+                return newNode
+            }
             cell = newNodeBuffer[0]
         } else {
             cell = newNode