-
Notifications
You must be signed in to change notification settings - Fork 0
/
treapiter.go
354 lines (319 loc) · 10.9 KB
/
treapiter.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
// Copyright (c) 2015-2016 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.
package treap
import "bytes"
// Iterator represents an iterator for forwards and backwards iteration over
// the contents of a treap (mutable or immutable).
type Iterator struct {
t *Mutable // Mutable treap iterator is associated with or nil
root *treapNode // Root node of treap iterator is associated with
node *treapNode // The node the iterator is positioned at
parents parentStack // The stack of parents needed to iterate
isNew bool // Whether the iterator has been positioned
seekKey []byte // Used to handle dynamic updates for mutable treap
startKey []byte // Used to limit the iterator to a range
limitKey []byte // Used to limit the iterator to a range
}
// limitIterator clears the current iterator node if it is outside of the range
// specified when the iterator was created. It returns whether the iterator is
// valid.
func (iter *Iterator) limitIterator() bool {
if iter.node == nil {
return false
}
node := iter.node
if iter.startKey != nil && bytes.Compare(node.key, iter.startKey) < 0 {
iter.node = nil
return false
}
if iter.limitKey != nil && bytes.Compare(node.key, iter.limitKey) >= 0 {
iter.node = nil
return false
}
return true
}
// seek moves the iterator based on the provided key and flags.
//
// When the exact match flag is set, the iterator will either be moved to first
// key in the treap that exactly matches the provided key, or the one
// before/after it depending on the greater flag.
//
// When the exact match flag is NOT set, the iterator will be moved to the first
// key in the treap before/after the provided key depending on the greater flag.
//
// In all cases, the limits specified when the iterator was created are
// respected.
func (iter *Iterator) seek(key []byte, exactMatch bool, greater bool) bool {
iter.node = nil
iter.parents = parentStack{}
var selectedNodeDepth int
for node := iter.root; node != nil; {
iter.parents.Push(node)
// Traverse left or right depending on the result of the
// comparison. Also, set the iterator to the node depending on
// the flags so the iterator is positioned properly when an
// exact match isn't found.
compareResult := bytes.Compare(key, node.key)
if compareResult < 0 {
if greater {
iter.node = node
selectedNodeDepth = iter.parents.Len() - 1
}
node = node.left
continue
}
if compareResult > 0 {
if !greater {
iter.node = node
selectedNodeDepth = iter.parents.Len() - 1
}
node = node.right
continue
}
// The key is an exact match. Set the iterator and return now
// when the exact match flag is set.
if exactMatch {
iter.node = node
iter.parents.Pop()
return iter.limitIterator()
}
// The key is an exact match, but the exact match is not set, so
// choose which direction to go based on whether the larger or
// smaller key was requested.
if greater {
node = node.right
} else {
node = node.left
}
}
// There was either no exact match or there was an exact match but the
// exact match flag was not set. In any case, the parent stack might
// need to be adjusted to only include the parents up to the selected
// node. Also, ensure the selected node's key does not exceed the
// allowed range of the iterator.
for i := iter.parents.Len(); i > selectedNodeDepth; i-- {
iter.parents.Pop()
}
return iter.limitIterator()
}
// First moves the iterator to the first key/value pair. When there is only a
// single key/value pair both First and Last will point to the same pair.
// Returns false if there are no key/value pairs.
func (iter *Iterator) First() bool {
// Seek the start key if the iterator was created with one. This will
// result in either an exact match, the first greater key, or an
// exhausted iterator if no such key exists.
iter.isNew = false
if iter.startKey != nil {
return iter.seek(iter.startKey, true, true)
}
// The smallest key is in the left-most node.
iter.parents = parentStack{}
for node := iter.root; node != nil; node = node.left {
if node.left == nil {
iter.node = node
return true
}
iter.parents.Push(node)
}
return false
}
// Last moves the iterator to the last key/value pair. When there is only a
// single key/value pair both First and Last will point to the same pair.
// Returns false if there are no key/value pairs.
func (iter *Iterator) Last() bool {
// Seek the limit key if the iterator was created with one. This will
// result in the first key smaller than the limit key, or an exhausted
// iterator if no such key exists.
iter.isNew = false
if iter.limitKey != nil {
return iter.seek(iter.limitKey, false, false)
}
// The highest key is in the right-most node.
iter.parents = parentStack{}
for node := iter.root; node != nil; node = node.right {
if node.right == nil {
iter.node = node
return true
}
iter.parents.Push(node)
}
return false
}
// Next moves the iterator to the next key/value pair and returns false when the
// iterator is exhausted. When invoked on a newly created iterator it will
// position the iterator at the first item.
func (iter *Iterator) Next() bool {
if iter.isNew {
return iter.First()
}
if iter.node == nil {
return false
}
// Reseek the previous key without allowing for an exact match if a
// force seek was requested. This results in the key greater than the
// previous one or an exhausted iterator if there is no such key.
if seekKey := iter.seekKey; seekKey != nil {
iter.seekKey = nil
return iter.seek(seekKey, false, true)
}
// When there is no right node walk the parents until the parent's right
// node is not equal to the previous child. This will be the next node.
if iter.node.right == nil {
parent := iter.parents.Pop()
for parent != nil && parent.right == iter.node {
iter.node = parent
parent = iter.parents.Pop()
}
iter.node = parent
return iter.limitIterator()
}
// There is a right node, so the next node is the left-most node down
// the right sub-tree.
iter.parents.Push(iter.node)
iter.node = iter.node.right
for node := iter.node.left; node != nil; node = node.left {
iter.parents.Push(iter.node)
iter.node = node
}
return iter.limitIterator()
}
// Prev moves the iterator to the previous key/value pair and returns false when
// the iterator is exhausted. When invoked on a newly created iterator it will
// position the iterator at the last item.
func (iter *Iterator) Prev() bool {
if iter.isNew {
return iter.Last()
}
if iter.node == nil {
return false
}
// Reseek the previous key without allowing for an exact match if a
// force seek was requested. This results in the key smaller than the
// previous one or an exhausted iterator if there is no such key.
if seekKey := iter.seekKey; seekKey != nil {
iter.seekKey = nil
return iter.seek(seekKey, false, false)
}
// When there is no left node walk the parents until the parent's left
// node is not equal to the previous child. This will be the previous
// node.
for iter.node.left == nil {
parent := iter.parents.Pop()
for parent != nil && parent.left == iter.node {
iter.node = parent
parent = iter.parents.Pop()
}
iter.node = parent
return iter.limitIterator()
}
// There is a left node, so the previous node is the right-most node
// down the left sub-tree.
iter.parents.Push(iter.node)
iter.node = iter.node.left
for node := iter.node.right; node != nil; node = node.right {
iter.parents.Push(iter.node)
iter.node = node
}
return iter.limitIterator()
}
// Seek moves the iterator to the first key/value pair with a key that is
// greater than or equal to the given key and returns true if successful.
func (iter *Iterator) Seek(key []byte) bool {
iter.isNew = false
return iter.seek(key, true, true)
}
// Key returns the key of the current key/value pair or nil when the iterator
// is exhausted. The caller should not modify the contents of the returned
// slice.
func (iter *Iterator) Key() []byte {
if iter.node == nil {
return nil
}
return iter.node.key
}
// Value returns the value of the current key/value pair or nil when the
// iterator is exhausted. The caller should not modify the contents of the
// returned slice.
func (iter *Iterator) Value() []byte {
if iter.node == nil {
return nil
}
return iter.node.value
}
// Valid indicates whether the iterator is positioned at a valid key/value pair.
// It will be considered invalid when the iterator is newly created or exhausted.
func (iter *Iterator) Valid() bool {
return iter.node != nil
}
// ForceReseek notifies the iterator that the underlying mutable treap has been
// updated, so the next call to Prev or Next needs to reseek in order to allow
// the iterator to continue working properly.
//
// NOTE: Calling this function when the iterator is associated with an immutable
// treap has no effect as you would expect.
func (iter *Iterator) ForceReseek() {
// Nothing to do when the iterator is associated with an immutable
// treap.
if iter.t == nil {
return
}
// Update the iterator root to the mutable treap root in case it
// changed.
iter.root = iter.t.root
// Set the seek key to the current node. This will force the Next/Prev
// functions to reseek, and thus properly reconstruct the iterator, on
// their next call.
if iter.node == nil {
iter.seekKey = nil
return
}
iter.seekKey = iter.node.key
}
// Iterator returns a new iterator for the mutable treap. The newly returned
// iterator is not pointing to a valid item until a call to one of the methods
// to position it is made.
//
// The start key and limit key parameters cause the iterator to be limited to
// a range of keys. The start key is inclusive and the limit key is exclusive.
// Either or both can be nil if the functionality is not desired.
//
// WARNING: The ForceSeek method must be called on the returned iterator if
// the treap is mutated. Failure to do so will cause the iterator to return
// unexpected keys and/or values.
//
// For example:
// iter := t.Iterator(nil, nil)
// for iter.Next() {
// if someCondition {
// t.Delete(iter.Key())
// iter.ForceReseek()
// }
// }
func (t *Mutable) Iterator(startKey, limitKey []byte) *Iterator {
iter := &Iterator{
t: t,
root: t.root,
isNew: true,
startKey: startKey,
limitKey: limitKey,
}
return iter
}
// Iterator returns a new iterator for the immutable treap. The newly returned
// iterator is not pointing to a valid item until a call to one of the methods
// to position it is made.
//
// The start key and limit key parameters cause the iterator to be limited to
// a range of keys. The start key is inclusive and the limit key is exclusive.
// Either or both can be nil if the functionality is not desired.
func (t *Immutable) Iterator(startKey, limitKey []byte) *Iterator {
iter := &Iterator{
root: t.root,
isNew: true,
startKey: startKey,
limitKey: limitKey,
}
return iter
}