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