Refactor TrieNode of PersistentHashMap

Author: ilya-g

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

@@ -95,7 +95,8 @@ internal class PersistentHashMap<K, V>(internal val node: TrieNode<K, V>,
     }
 
     internal companion object {
-        private val EMPTY = PersistentHashMap(TrieNode.EMPTY, 0)
+        private val EMPTY = PersistentHashMap<Nothing, Nothing>(TrieNode.EMPTY, 0)
+        @Suppress("UNCHECKED_CAST")
         internal fun <K, V> emptyOf(): PersistentHashMap<K, V> = EMPTY as PersistentHashMap<K, V>
     }
 }

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

@@ -72,28 +72,28 @@ internal abstract class PersistentHashMapBuilderBaseIterator<K, V, T>(private va
 
     private fun resetPath(keyHash: Int, node: TrieNode<*, *>, key: K, pathIndex: Int) {
         val shift = pathIndex * LOG_MAX_BRANCHING_FACTOR
-        val keyPosition = 1 shl ((keyHash shr shift) and MAX_BRANCHING_FACTOR_MINUS_ONE)
+        val keyPositionMask = 1 shl indexSegment(keyHash, shift)
 
-        if (node.hasDataAt(keyPosition)) { // key is directly in buffer
-            val keyIndex = node.keyDataIndex(keyPosition)
+        if (node.hasEntryAt(keyPositionMask)) { // key is directly in buffer
+            val keyIndex = node.entryKeyIndex(keyPositionMask)
 
 //            assert(node.keyAtIndex(keyIndex) == key)
 
-            path[pathIndex].reset(node.buffer, 2 * Integer.bitCount(node.dataMap), keyIndex)
+            path[pathIndex].reset(node.buffer, ENTRY_SIZE * node.entryCount(), keyIndex)
             return
         }
 
 //        assert(node.hasNodeAt(keyPosition)) // key is in node
 
-        val nodeIndex = node.keyNodeIndex(keyPosition)
+        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, 2 * Integer.bitCount(node.dataMap), nodeIndex)
+            path[pathIndex].reset(node.buffer, ENTRY_SIZE * node.entryCount(), nodeIndex)
             resetPath(keyHash, targetNode, key, pathIndex + 1)
         }
     }

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

@@ -4,6 +4,7 @@ internal const val TRIE_MAX_HEIGHT = 7
 
 internal abstract class TrieNodeBaseIterator<out K, out V, out T> : Iterator<T> {
     protected var buffer = emptyArray<Any?>()
+        private set
     private var dataSize = 0
     protected var index = 0
 
@@ -90,7 +91,7 @@ internal abstract class PersistentHashMapBaseIterator<K, V, T>(node: TrieNode<K,
     private var hasNext = true
 
     init {
-        path[0].reset(node.buffer, 2 * Integer.bitCount(node.dataMap))
+        path[0].reset(node.buffer, ENTRY_SIZE * node.entryCount())
         pathLastIndex = 0
         ensureNextEntryIsReady()
     }
@@ -104,7 +105,7 @@ internal abstract class PersistentHashMapBaseIterator<K, V, T>(node: TrieNode<K,
             if (pathIndex == TRIE_MAX_HEIGHT - 1) {     // collision
                 path[pathIndex + 1].reset(node.buffer, node.buffer.size)
             } else {
-                path[pathIndex + 1].reset(node.buffer, 2 * Integer.bitCount(node.dataMap))
+                path[pathIndex + 1].reset(node.buffer, ENTRY_SIZE * node.entryCount())
             }
             return moveToNextNodeWithData(pathIndex + 1)
         }