Do note create node containing a single other node

Author: Abduqodiri Qurbonzoda

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

@@ -183,12 +183,19 @@ 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)
         // TODO: check how this changes affect `put` and non-collision `remove` operations performance.
 
         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)
@@ -199,11 +206,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])
 

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

@@ -98,29 +98,41 @@ 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)
+        // TODO: check how this changes affect `add` and non-collision `remove` operations performance.
 
-        val newBuffer = buffer.copyOf()
+        val cell: Any?
 
-        // 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