1 package treap
2 3 import (
4 "bytes"
5 "math/rand"
6 )
7 8 // Mutable represents a treap data structure which is used to hold ordered key/value pairs using a combination of binary
9 // search tree and heap semantics. It is a self-organizing and randomized data structure that doesn't require complex
10 // operations to maintain balance. Search, insert, and delete operations are all O(log n).
11 type Mutable struct {
12 root *treapNode
13 count int
14 // totalSize is the best estimate of the total size of of all data in the treap including the keys, values, and node txsizes.
15 totalSize uint64
16 }
17 18 // Len returns the number of items stored in the treap.
19 func (t *Mutable) Len() int {
20 return t.count
21 }
22 23 // Size returns a best estimate of the total number of bytes the treap is consuming including all of the fields used to
24 // represent the nodes as well as the size of the keys and values. Shared values are not detected, so the returned size
25 // assumes each value is pointing to different memory.
26 func (t *Mutable) Size() uint64 {
27 return t.totalSize
28 }
29 30 // get returns the treap node that contains the passed key and its parent. When the found node is the root of the tree,
31 // the parent will be nil. When the key does not exist, both the node and the parent will be nil.
32 func (t *Mutable) get(key []byte) (*treapNode, *treapNode) {
33 var parent *treapNode
34 for node := t.root; node != nil; {
35 // Traverse left or right depending on the result of the comparison.
36 compareResult := bytes.Compare(key, node.key)
37 if compareResult < 0 {
38 parent = node
39 node = node.left
40 continue
41 }
42 if compareResult > 0 {
43 parent = node
44 node = node.right
45 continue
46 }
47 // The key exists.
48 return node, parent
49 }
50 // A nil node was reached which means the key does not exist.
51 return nil, nil
52 }
53 54 // Has returns whether or not the passed key exists.
55 func (t *Mutable) Has(key []byte) bool {
56 if node, _ := t.get(key); node != nil {
57 return true
58 }
59 return false
60 }
61 62 // Get returns the value for the passed key. The function will return nil when the key does not exist.
63 func (t *Mutable) Get(key []byte) []byte {
64 if node, _ := t.get(key); node != nil {
65 return node.value
66 }
67 return nil
68 }
69 70 // relinkGrandparent relinks the node into the treap after it has been rotated by changing the passed grandparent's left
71 // or right pointer, depending on where the old parent was, to point at the passed node. Otherwise, when there is no
72 // grandparent, it means the node is now the root of the tree, so update it accordingly.
73 func (t *Mutable) relinkGrandparent(node, parent, grandparent *treapNode) {
74 // The node is now the root of the tree when there is no grandparent.
75 if grandparent == nil {
76 t.root = node
77 return
78 }
79 // Relink the grandparent's left or right pointer based on which side the old parent was.
80 if grandparent.left == parent {
81 grandparent.left = node
82 } else {
83 grandparent.right = node
84 }
85 }
86 87 // Put inserts the passed key/value pair.
88 func (t *Mutable) Put(key, value []byte) {
89 // Use an empty byte slice for the value when none was provided. This ultimately allows key existence to be
90 // determined from the value since an empty byte slice is distinguishable from nil.
91 if value == nil {
92 value = emptySlice
93 }
94 // The node is the root of the tree if there isn't already one.
95 if t.root == nil {
96 node := newTreapNode(key, value, rand.Int())
97 t.count = 1
98 t.totalSize = nodeSize(node)
99 t.root = node
100 return
101 }
102 // Find the binary tree insertion point and construct a list of parents while doing so. When the key matches an
103 // entry already in the treap, just update its value and return.
104 var parents parentStack
105 var compareResult int
106 for node := t.root; node != nil; {
107 parents.Push(node)
108 compareResult = bytes.Compare(key, node.key)
109 if compareResult < 0 {
110 node = node.left
111 continue
112 }
113 if compareResult > 0 {
114 node = node.right
115 continue
116 }
117 // The key already exists, so update its value.
118 t.totalSize -= uint64(len(node.value))
119 t.totalSize += uint64(len(value))
120 node.value = value
121 return
122 }
123 // Link the new node into the binary tree in the correct position.
124 node := newTreapNode(key, value, rand.Int())
125 t.count++
126 t.totalSize += nodeSize(node)
127 parent := parents.At(0)
128 if compareResult < 0 {
129 parent.left = node
130 } else {
131 parent.right = node
132 }
133 // Perform any rotations needed to maintain the min-heap.
134 for parents.Len() > 0 {
135 // There is nothing left to do when the node's priority is greater than or equal to its parent's priority.
136 parent = parents.Pop()
137 if node.priority >= parent.priority {
138 break
139 }
140 // Perform a right rotation if the node is on the left side or a left rotation if the node is on the right side.
141 if parent.left == node {
142 node.right, parent.left = parent, node.right
143 } else {
144 node.left, parent.right = parent, node.left
145 }
146 t.relinkGrandparent(node, parent, parents.At(0))
147 }
148 }
149 150 // Delete removes the passed key if it exists.
151 func (t *Mutable) Delete(key []byte) {
152 // Find the node for the key along with its parent. There is nothing to do if the key does not exist.
153 node, parent := t.get(key)
154 if node == nil {
155 return
156 }
157 // When the only node in the tree is the root node and it is the one being deleted, there is nothing else to do besides removing it.
158 if parent == nil && node.left == nil && node.right == nil {
159 t.root = nil
160 t.count = 0
161 t.totalSize = 0
162 return
163 }
164 // Perform rotations to move the node to delete to a leaf position while maintaining the min-heap.
165 var isLeft bool
166 var child *treapNode
167 for node.left != nil || node.right != nil {
168 // Choose the child with the higher priority.
169 if node.left == nil {
170 child = node.right
171 isLeft = false
172 } else if node.right == nil {
173 child = node.left
174 isLeft = true
175 } else if node.left.priority >= node.right.priority {
176 child = node.left
177 isLeft = true
178 } else {
179 child = node.right
180 isLeft = false
181 }
182 // Rotate left or right depending on which side the child node is on. This has the effect of moving the node to
183 // delete towards the bottom of the tree while maintaining the min-heap.
184 if isLeft {
185 child.right, node.left = node, child.right
186 } else {
187 child.left, node.right = node, child.left
188 }
189 t.relinkGrandparent(child, node, parent)
190 // The parent for the node to delete is now what was previously its child.
191 parent = child
192 }
193 // Delete the node, which is now a leaf node, by disconnecting it from its parent.
194 if parent != nil {
195 if parent.right == node {
196 parent.right = nil
197 } else {
198 parent.left = nil
199 }
200 t.count--
201 t.totalSize -= nodeSize(node)
202 }
203 }
204 205 // ForEach invokes the passed function with every key/value pair in the treap in ascending order.
206 func (t *Mutable) ForEach(fn func(k, v []byte) bool) {
207 // Add the root node and all children to the left of it to the list of nodes to traverse and loop until they, and
208 // all of their child nodes, been traversed.
209 var parents parentStack
210 for node := t.root; node != nil; node = node.left {
211 parents.Push(node)
212 }
213 for parents.Len() > 0 {
214 node := parents.Pop()
215 if !fn(node.key, node.value) {
216 return
217 }
218 // Extend the nodes to traverse by all children to the left of the current node's right child.
219 for node = node.right; node != nil; node = node.left {
220 parents.Push(node)
221 }
222 }
223 }
224 225 // Reset efficiently removes all items in the treap.
226 func (t *Mutable) Reset() {
227 t.count = 0
228 t.totalSize = 0
229 t.root = nil
230 }
231 232 // NewMutable returns a new empty mutable treap ready for use. See the documentation for the Mutable structure for more
233 // details.
234 func NewMutable() *Mutable {
235 return &Mutable{}
236 }
237