Skip to content

Commit fdd73d6

Browse files
author
Abduqodiri Qurbonzoda
committed
Implement map TrieNode canonicalization
1 parent dca8d17 commit fdd73d6

File tree

1 file changed

+52
-51
lines changed
  • kotlinx-collections-immutable/src/main/kotlin/kotlinx/collections/immutable/implementations/immutableMap

1 file changed

+52
-51
lines changed

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

Lines changed: 52 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -52,20 +52,22 @@ private fun Array<Any?>.replaceEntryWithNode(keyIndex: Int, nodeIndex: Int, newN
5252
return newBuffer
5353
}
5454

55+
private fun <K, V> Array<Any?>.replaceNodeWithEntry(nodeIndex: Int, keyIndex: Int, key: K, value: V): Array<Any?> {
56+
val newBuffer = this.copyOf(this.size + 1)
57+
newBuffer.copyInto(newBuffer, nodeIndex + 2, nodeIndex + 1, this.size)
58+
newBuffer.copyInto(newBuffer, keyIndex + 2, keyIndex , nodeIndex)
59+
newBuffer[keyIndex] = key
60+
newBuffer[keyIndex + 1] = value
61+
return newBuffer
62+
}
63+
5564
private fun Array<Any?>.removeEntryAtIndex(keyIndex: Int): Array<Any?> {
5665
val newBuffer = arrayOfNulls<Any?>(this.size - ENTRY_SIZE)
5766
this.copyInto(newBuffer, endIndex = keyIndex)
5867
this.copyInto(newBuffer, keyIndex, startIndex = keyIndex + ENTRY_SIZE, endIndex = this.size)
5968
return newBuffer
6069
}
6170

62-
private fun Array<Any?>.removeNodeAtIndex(nodeIndex: Int): Array<Any?> {
63-
val newBuffer = arrayOfNulls<Any?>(this.size - 1)
64-
this.copyInto(newBuffer, endIndex = nodeIndex)
65-
this.copyInto(newBuffer, nodeIndex, startIndex = nodeIndex + 1, endIndex = this.size)
66-
return newBuffer
67-
}
68-
6971

7072

7173
internal class TrieNode<K, V>(
@@ -181,17 +183,40 @@ internal class TrieNode<K, V>(
181183
return TrieNode(dataMap, nodeMap, newBuffer, mutator.ownership)
182184
}
183185

184-
private fun updateNodeAtIndex(nodeIndex: Int, newNode: TrieNode<K, V>): TrieNode<K, V> {
186+
private fun updateNodeAtIndex(nodeIndex: Int, positionMask: Int, newNode: TrieNode<K, V>): TrieNode<K, V> {
185187
// assert(buffer[nodeIndex] !== newNode)
188+
// TODO: check how this changes affect `put` and non-collision `remove` operations performance.
186189

187-
val newBuffer = buffer.copyOf()
190+
val newNodeBuffer = newNode.buffer
191+
if (newNodeBuffer.size == 2 && newNode.nodeMap == 0) {
192+
val keyIndex = entryKeyIndex(positionMask)
193+
val newBuffer = buffer.replaceNodeWithEntry(nodeIndex, keyIndex, newNodeBuffer[0], newNodeBuffer[1])
194+
return TrieNode(dataMap xor positionMask, nodeMap xor positionMask, newBuffer)
195+
}
196+
197+
val newBuffer = buffer.copyOf(buffer.size)
188198
newBuffer[nodeIndex] = newNode
189199
return TrieNode(dataMap, nodeMap, newBuffer)
190200
}
191201

192-
private fun mutableUpdateNodeAtIndex(nodeIndex: Int, newNode: TrieNode<K, V>, owner: MutabilityOwnership): TrieNode<K, V> {
202+
private fun mutableUpdateNodeAtIndex(nodeIndex: Int, positionMask: Int, newNode: TrieNode<K, V>, owner: MutabilityOwnership): TrieNode<K, V> {
193203
// assert(buffer[nodeIndex] !== newNode)
194204

205+
val newNodeBuffer = newNode.buffer
206+
if (newNodeBuffer.size == 2 && newNode.nodeMap == 0) {
207+
val keyIndex = entryKeyIndex(positionMask)
208+
val newBuffer = buffer.replaceNodeWithEntry(nodeIndex, keyIndex, newNodeBuffer[0], newNodeBuffer[1])
209+
210+
if (ownedBy === owner) {
211+
buffer = newBuffer
212+
dataMap = dataMap xor positionMask
213+
nodeMap = nodeMap xor positionMask
214+
return this
215+
}
216+
217+
return TrieNode(dataMap xor positionMask, nodeMap xor positionMask, newBuffer)
218+
}
219+
195220
if (ownedBy === owner) {
196221
buffer[nodeIndex] = newNode
197222
return this
@@ -282,6 +307,7 @@ internal class TrieNode<K, V>(
282307

283308
private fun removeEntryAtIndex(keyIndex: Int, positionMask: Int): TrieNode<K, V>? {
284309
// assert(hasEntryAt(positionMask))
310+
// It is possible only when this node is the root node
285311
if (buffer.size == ENTRY_SIZE) return null
286312

287313
val newBuffer = buffer.removeEntryAtIndex(keyIndex)
@@ -323,27 +349,6 @@ internal class TrieNode<K, V>(
323349
return makeCollisionNode(newBuffer, mutator.ownership)
324350
}
325351

326-
private fun removeNodeAtIndex(nodeIndex: Int, positionMask: Int): TrieNode<K, V>? {
327-
// assert(hasNodeAt(positionMask))
328-
if (buffer.size == 1) return null
329-
330-
val newBuffer = buffer.removeNodeAtIndex(nodeIndex)
331-
return TrieNode(dataMap, nodeMap xor positionMask, newBuffer)
332-
}
333-
334-
private fun mutableRemoveNodeAtIndex(nodeIndex: Int, positionMask: Int, owner: MutabilityOwnership): TrieNode<K, V>? {
335-
// assert(hasNodeAt(positionMask))
336-
if (buffer.size == 1) return null
337-
338-
if (ownedBy === owner) {
339-
buffer = buffer.removeNodeAtIndex(nodeIndex)
340-
nodeMap = nodeMap xor positionMask
341-
return this
342-
}
343-
val newBuffer = buffer.removeNodeAtIndex(nodeIndex)
344-
return TrieNode(dataMap, nodeMap xor positionMask, newBuffer, owner)
345-
}
346-
347352
private fun collisionContainsKey(key: K): Boolean {
348353
val collisionHash = keyAtIndex(0).hashCode()
349354
val keyHash = key.hashCode()
@@ -559,7 +564,7 @@ internal class TrieNode<K, V>(
559564

560565
val targetNode = nodeAtIndex(nodeIndex)
561566
val putResult = targetNode.put(keyHash, key, value, shift + LOG_MAX_BRANCHING_FACTOR) ?: return null
562-
return putResult.replaceNode { node -> updateNodeAtIndex(nodeIndex, node) }
567+
return putResult.replaceNode { node -> updateNodeAtIndex(nodeIndex, keyPositionMask, node) }
563568
}
564569

565570
// no entry at this key hash segment
@@ -595,7 +600,7 @@ internal class TrieNode<K, V>(
595600
if (targetNode === newNode) {
596601
return this
597602
}
598-
return mutableUpdateNodeAtIndex(nodeIndex, newNode, mutator.ownership)
603+
return mutableUpdateNodeAtIndex(nodeIndex, keyPositionMask, newNode, mutator.ownership)
599604
}
600605

601606
// key is absent
@@ -623,11 +628,9 @@ internal class TrieNode<K, V>(
623628

624629
val targetNode = nodeAtIndex(nodeIndex)
625630
val newNode = targetNode.remove(keyHash, key, shift + LOG_MAX_BRANCHING_FACTOR)
626-
return when {
627-
targetNode === newNode -> this
628-
newNode == null -> removeNodeAtIndex(nodeIndex, keyPositionMask)
629-
else -> updateNodeAtIndex(nodeIndex, newNode)
630-
}
631+
checkNotNull(newNode)
632+
if (targetNode === newNode) return this
633+
return updateNodeAtIndex(nodeIndex, keyPositionMask, newNode)
631634
}
632635

633636
// key is absent
@@ -654,11 +657,11 @@ internal class TrieNode<K, V>(
654657

655658
val targetNode = nodeAtIndex(nodeIndex)
656659
val newNode = targetNode.mutableRemove(keyHash, key, shift + LOG_MAX_BRANCHING_FACTOR, mutator)
657-
return when {
658-
targetNode === newNode -> this
659-
newNode == null -> mutableRemoveNodeAtIndex(nodeIndex, keyPositionMask, mutator.ownership)
660-
else -> mutableUpdateNodeAtIndex(nodeIndex, newNode, mutator.ownership)
660+
checkNotNull(newNode)
661+
if (ownedBy === mutator.ownership || targetNode !== newNode) {
662+
return mutableUpdateNodeAtIndex(nodeIndex, keyPositionMask, newNode, mutator.ownership)
661663
}
664+
return this
662665
}
663666

664667
// key is absent
@@ -685,11 +688,9 @@ internal class TrieNode<K, V>(
685688

686689
val targetNode = nodeAtIndex(nodeIndex)
687690
val newNode = targetNode.remove(keyHash, key, value, shift + LOG_MAX_BRANCHING_FACTOR)
688-
return when {
689-
targetNode === newNode -> this
690-
newNode == null -> removeNodeAtIndex(nodeIndex, keyPositionMask)
691-
else -> updateNodeAtIndex(nodeIndex, newNode)
692-
}
691+
checkNotNull(newNode)
692+
if (targetNode === newNode) return this
693+
return updateNodeAtIndex(nodeIndex, keyPositionMask, newNode)
693694
}
694695

695696
// key is absent
@@ -716,11 +717,11 @@ internal class TrieNode<K, V>(
716717

717718
val targetNode = nodeAtIndex(nodeIndex)
718719
val newNode = targetNode.mutableRemove(keyHash, key, value, shift + LOG_MAX_BRANCHING_FACTOR, mutator)
719-
return when {
720-
targetNode === newNode -> this
721-
newNode == null -> mutableRemoveNodeAtIndex(nodeIndex, keyPositionMask, mutator.ownership)
722-
else -> mutableUpdateNodeAtIndex(nodeIndex, newNode, mutator.ownership)
720+
checkNotNull(newNode)
721+
if (ownedBy === mutator.ownership || targetNode !== newNode) {
722+
return mutableUpdateNodeAtIndex(nodeIndex, keyPositionMask, newNode, mutator.ownership)
723723
}
724+
return this
724725
}
725726

726727
// key is absent

0 commit comments

Comments
 (0)