1 package treap
2 3 import (
4 "bytes"
5 "math/rand"
6 )
7 8 // cloneTreapNode returns a shallow copy of the passed node.
9 func cloneTreapNode(node *treapNode) *treapNode {
10 return &treapNode{
11 key: node.key,
12 value: node.value,
13 priority: node.priority,
14 left: node.left,
15 right: node.right,
16 }
17 }
18 19 // Immutable represents a treap data structure which is used to hold ordered key/value pairs using a combination of
20 // binary search tree and heap semantics. It is a self-organizing and randomized data structure that doesn't require
21 // complex operations to maintain balance. Search, insert, and delete operations are all O(log n).
22 //
23 // In addition, it provides O(1) snapshots for multi-version concurrency control (MVCC). All operations which result in
24 // modifying the treap return a new version of the treap with only the modified nodes updated. All unmodified nodes are
25 // shared with the previous version.
26 //
27 // This is extremely useful in concurrent applications since the caller only has to atomically replace the treap pointer
28 // with the newly returned version after performing any mutations. All readers can simply use their existing pointer as
29 // a snapshot since the treap it points to is immutable.
30 //
31 // This effectively provides O(1) snapshot capability with efficient memory usage characteristics since the old nodes
32 // only remain allocated until there are no longer any references to them.
33 type Immutable struct {
34 root *treapNode
35 count int
36 // totalSize is the best estimate of the total size of of all data in the treap including the keys, values, and node
37 // txsizes.
38 totalSize uint64
39 }
40 41 // newImmutable returns a new immutable treap given the passed parameters.
42 func newImmutable(root *treapNode, count int, totalSize uint64) *Immutable {
43 return &Immutable{root: root, count: count, totalSize: totalSize}
44 }
45 46 // Len returns the number of items stored in the treap.
47 func (t *Immutable) Len() int {
48 return t.count
49 }
50 51 // Size returns a best estimate of the total number of bytes the treap is consuming including all of the fields used to
52 // represent the nodes as well as the size of the keys and values. Shared values are not detected, so the returned size
53 // assumes each value is pointing to different memory.
54 func (t *Immutable) Size() uint64 {
55 return t.totalSize
56 }
57 58 // get returns the treap node that contains the passed key. It will return nil when the key does not exist.
59 func (t *Immutable) get(key []byte) *treapNode {
60 for node := t.root; node != nil; {
61 // Traverse left or right depending on the result of the comparison.
62 compareResult := bytes.Compare(key, node.key)
63 if compareResult < 0 {
64 node = node.left
65 continue
66 }
67 if compareResult > 0 {
68 node = node.right
69 continue
70 }
71 // The key exists.
72 return node
73 }
74 // A nil node was reached which means the key does not exist.
75 return nil
76 }
77 78 // Has returns whether or not the passed key exists.
79 func (t *Immutable) Has(key []byte) bool {
80 if node := t.get(key); node != nil {
81 return true
82 }
83 return false
84 }
85 86 // Get returns the value for the passed key. The function will return nil when the key does not exist.
87 func (t *Immutable) Get(key []byte) []byte {
88 if node := t.get(key); node != nil {
89 return node.value
90 }
91 return nil
92 }
93 94 // Put inserts the passed key/value pair.
95 func (t *Immutable) Put(key, value []byte) *Immutable {
96 // Use an empty byte slice for the value when none was provided. This ultimately allows key existence to be
97 // determined from the value since an empty byte slice is distinguishable from nil.
98 if value == nil {
99 value = emptySlice
100 }
101 // The node is the root of the tree if there isn't already one.
102 if t.root == nil {
103 root := newTreapNode(key, value, rand.Int())
104 return newImmutable(root, 1, nodeSize(root))
105 }
106 // Find the binary tree insertion point and construct a replaced list of parents while doing so. This is done
107 // because this is an immutable data structure so regardless of where in the treap the new key/value pair ends up,
108 // all ancestors up to and including the root need to be replaced.
109 //
110 // When the key matches an entry already in the treap, replace the node with a new one that has the new value set
111 // and return.
112 var parents parentStack
113 var compareResult int
114 for node := t.root; node != nil; {
115 // Clone the node and link its parent to it if needed.
116 nodeCopy := cloneTreapNode(node)
117 if oldParent := parents.At(0); oldParent != nil {
118 if oldParent.left == node {
119 oldParent.left = nodeCopy
120 } else {
121 oldParent.right = nodeCopy
122 }
123 }
124 parents.Push(nodeCopy)
125 // Traverse left or right depending on the result of comparing the keys.
126 compareResult = bytes.Compare(key, node.key)
127 if compareResult < 0 {
128 node = node.left
129 continue
130 }
131 if compareResult > 0 {
132 node = node.right
133 continue
134 }
135 // The key already exists, so update its value.
136 nodeCopy.value = value
137 // Return new immutable treap with the replaced node and ancestors up to and including the root of the tree.
138 newRoot := parents.At(parents.Len() - 1)
139 newTotalSize := t.totalSize - uint64(len(node.value)) +
140 uint64(len(value))
141 return newImmutable(newRoot, t.count, newTotalSize)
142 }
143 // Link the new node into the binary tree in the correct position.
144 node := newTreapNode(key, value, rand.Int())
145 parent := parents.At(0)
146 if compareResult < 0 {
147 parent.left = node
148 } else {
149 parent.right = node
150 }
151 // Perform any rotations needed to maintain the min-heap and replace the ancestors up to and including the tree
152 // root.
153 newRoot := parents.At(parents.Len() - 1)
154 for parents.Len() > 0 {
155 // There is nothing left to do when the node's priority is greater than or equal to its parent's priority.
156 parent = parents.Pop()
157 if node.priority >= parent.priority {
158 break
159 }
160 // Perform a right rotation if the node is on the left side or a left rotation if the node is on the right side.
161 if parent.left == node {
162 node.right, parent.left = parent, node.right
163 } else {
164 node.left, parent.right = parent, node.left
165 }
166 // Either set the new root of the tree when there is no grandparent or relink the grandparent to the node based
167 // on which side the old parent the node is replacing was on.
168 grandparent := parents.At(0)
169 if grandparent == nil {
170 newRoot = node
171 } else if grandparent.left == parent {
172 grandparent.left = node
173 } else {
174 grandparent.right = node
175 }
176 }
177 return newImmutable(newRoot, t.count+1, t.totalSize+nodeSize(node))
178 }
179 180 // Delete removes the passed key from the treap and returns the resulting treap if it exists. The original immutable
181 // treap is returned if the key does not exist.
182 func (t *Immutable) Delete(key []byte) *Immutable {
183 // Find the node for the key while constructing a list of parents while doing so.
184 var parents parentStack
185 var delNode *treapNode
186 for node := t.root; node != nil; {
187 parents.Push(node)
188 // Traverse left or right depending on the result of the comparison.
189 compareResult := bytes.Compare(key, node.key)
190 if compareResult < 0 {
191 node = node.left
192 continue
193 }
194 if compareResult > 0 {
195 node = node.right
196 continue
197 }
198 // The key exists.
199 delNode = node
200 break
201 }
202 // There is nothing to do if the key does not exist.
203 if delNode == nil {
204 return t
205 }
206 // When the only node in the tree is the root node and it is the one being deleted, there is nothing else to do
207 // besides removing it.
208 parent := parents.At(1)
209 if parent == nil && delNode.left == nil && delNode.right == nil {
210 return newImmutable(nil, 0, 0)
211 }
212 // Construct a replaced list of parents and the node to delete itself. This is done because this is an immutable
213 // data structure and therefore all ancestors of the node that will be deleted, up to and including the root, need
214 // to be replaced.
215 var newParents parentStack
216 for i := parents.Len(); i > 0; i-- {
217 node := parents.At(i - 1)
218 nodeCopy := cloneTreapNode(node)
219 if oldParent := newParents.At(0); oldParent != nil {
220 if oldParent.left == node {
221 oldParent.left = nodeCopy
222 } else {
223 oldParent.right = nodeCopy
224 }
225 }
226 newParents.Push(nodeCopy)
227 }
228 delNode = newParents.Pop()
229 parent = newParents.At(0)
230 // Perform rotations to move the node to delete to a leaf position while maintaining the min-heap while replacing
231 // the modified children.
232 var child *treapNode
233 newRoot := newParents.At(newParents.Len() - 1)
234 for delNode.left != nil || delNode.right != nil {
235 // Choose the child with the higher priority.
236 var isLeft bool
237 if delNode.left == nil {
238 child = delNode.right
239 } else if delNode.right == nil {
240 child = delNode.left
241 isLeft = true
242 } else if delNode.left.priority >= delNode.right.priority {
243 child = delNode.left
244 isLeft = true
245 } else {
246 child = delNode.right
247 }
248 // Rotate left or right depending on which side the child node is on. This has the effect of moving the node to
249 // delete towards the bottom of the tree while maintaining the min-heap.
250 child = cloneTreapNode(child)
251 if isLeft {
252 child.right, delNode.left = delNode, child.right
253 } else {
254 child.left, delNode.right = delNode, child.left
255 }
256 // Either set the new root of the tree when there is no grandparent or relink the grandparent to the node based
257 // on which side the old parent the node is replacing was on. Since the node to be deleted was just moved down a
258 // level, the new grandparent is now the current parent and the new parent is the current child.
259 if parent == nil {
260 newRoot = child
261 } else if parent.left == delNode {
262 parent.left = child
263 } else {
264 parent.right = child
265 }
266 // The parent for the node to delete is now what was previously its child.
267 parent = child
268 }
269 // Delete the node, which is now a leaf node, by disconnecting it from its parent.
270 if parent.right == delNode {
271 parent.right = nil
272 } else {
273 parent.left = nil
274 }
275 return newImmutable(newRoot, t.count-1, t.totalSize-nodeSize(delNode))
276 }
277 278 // ForEach invokes the passed function with every key/value pair in the treap in ascending order.
279 func (t *Immutable) ForEach(fn func(k, v []byte) bool) {
280 // Add the root node and all children to the left of it to the list of nodes to traverse and loop until they, and
281 // all of their child nodes, have been traversed.
282 var parents parentStack
283 for node := t.root; node != nil; node = node.left {
284 parents.Push(node)
285 }
286 for parents.Len() > 0 {
287 node := parents.Pop()
288 if !fn(node.key, node.value) {
289 return
290 }
291 // Extend the nodes to traverse by all children to the left of the current node's right child.
292 for node = node.right; node != nil; node = node.left {
293 parents.Push(node)
294 }
295 }
296 }
297 298 // NewImmutable returns a new empty immutable treap ready for use. See the documentation for the Immutable structure for
299 // more details.
300 func NewImmutable() *Immutable {
301 return &Immutable{}
302 }
303