1 package treap
2 3 import "bytes"
4 5 // Iterator represents an iterator for forwards and backwards iteration over the contents of a treap (mutable or immutable).
6 type Iterator struct {
7 t *Mutable // Mutable treap iterator is associated with or nil
8 root *treapNode // Root node of treap iterator is associated with
9 node *treapNode // The node the iterator is positioned at
10 parents parentStack // The stack of parents needed to iterate
11 isNew bool // Whether the iterator has been positioned
12 seekKey []byte // Used to handle dynamic updates for mutable treap
13 startKey []byte // Used to limit the iterator to a range
14 limitKey []byte // Used to limit the iterator to a range
15 }
16 17 // limitIterator clears the current iterator node if it is outside of the range specified when the iterator was created.
18 // It returns whether the iterator is valid.
19 func (iter *Iterator) limitIterator() bool {
20 if iter.node == nil {
21 return false
22 }
23 node := iter.node
24 if iter.startKey != nil && bytes.Compare(node.key, iter.startKey) < 0 {
25 iter.node = nil
26 return false
27 }
28 if iter.limitKey != nil && bytes.Compare(node.key, iter.limitKey) >= 0 {
29 iter.node = nil
30 return false
31 }
32 return true
33 }
34 35 // seek moves the iterator based on the provided key and flags. When the exact match flag is set, the iterator will
36 // either be moved to first key in the treap that exactly matches the provided key, or the one before/after it depending
37 // on the greater flag. When the exact match flag is NOT set, the iterator will be moved to the first key in the treap
38 // before/after the provided key depending on the greater flag. In all cases, the limits specified when the iterator was
39 // created are respected.
40 func (iter *Iterator) seek(key []byte, exactMatch bool, greater bool) bool {
41 iter.node = nil
42 iter.parents = parentStack{}
43 var selectedNodeDepth int
44 for node := iter.root; node != nil; {
45 iter.parents.Push(node)
46 // Traverse left or right depending on the result of the comparison. Also, set the iterator to the node
47 // depending on the flags so the iterator is positioned properly when an exact match isn't found.
48 compareResult := bytes.Compare(key, node.key)
49 if compareResult < 0 {
50 if greater {
51 iter.node = node
52 selectedNodeDepth = iter.parents.Len() - 1
53 }
54 node = node.left
55 continue
56 }
57 if compareResult > 0 {
58 if !greater {
59 iter.node = node
60 selectedNodeDepth = iter.parents.Len() - 1
61 }
62 node = node.right
63 continue
64 }
65 // The key is an exact match. Set the iterator and return now when the exact match flag is set.
66 if exactMatch {
67 iter.node = node
68 iter.parents.Pop()
69 return iter.limitIterator()
70 }
71 // The key is an exact match, but the exact match is not set, so choose which direction to go based on whether
72 // the larger or smaller key was requested.
73 if greater {
74 node = node.right
75 } else {
76 node = node.left
77 }
78 }
79 // There was either no exact match or there was an exact match but the exact match flag was not set. In any case,
80 // the parent stack might need to be adjusted to only include the parents up to the selected node. Also, ensure the
81 // selected node's key does not exceed the allowed range of the iterator.
82 for i := iter.parents.Len(); i > selectedNodeDepth; i-- {
83 iter.parents.Pop()
84 }
85 return iter.limitIterator()
86 }
87 88 // First moves the iterator to the first key/value pair. When there is only a single key/value pair both First and Last
89 // will point to the same pair. Returns false if there are no key/value pairs.
90 func (iter *Iterator) First() bool {
91 // Seek the start key if the iterator was created with one. This will result in either an exact match, the first
92 // greater key, or an exhausted iterator if no such key exists.
93 iter.isNew = false
94 if iter.startKey != nil {
95 return iter.seek(iter.startKey, true, true)
96 }
97 // The smallest key is in the left-most node.
98 iter.parents = parentStack{}
99 for node := iter.root; node != nil; node = node.left {
100 if node.left == nil {
101 iter.node = node
102 return true
103 }
104 iter.parents.Push(node)
105 }
106 return false
107 }
108 109 // Last moves the iterator to the last key/value pair. When there is only a single key/value pair both First and Last
110 // will point to the same pair. Returns false if there are no key/value pairs.
111 func (iter *Iterator) Last() bool {
112 // Seek the limit key if the iterator was created with one. This will result in the first key smaller than the limit
113 // key, or an exhausted iterator if no such key exists.
114 iter.isNew = false
115 if iter.limitKey != nil {
116 return iter.seek(iter.limitKey, false, false)
117 }
118 // The highest key is in the right-most node.
119 iter.parents = parentStack{}
120 for node := iter.root; node != nil; node = node.right {
121 if node.right == nil {
122 iter.node = node
123 return true
124 }
125 iter.parents.Push(node)
126 }
127 return false
128 }
129 130 // Next moves the iterator to the next key/value pair and returns false when the iterator is exhausted. When invoked on
131 // a newly created iterator it will position the iterator at the first item.
132 func (iter *Iterator) Next() bool {
133 if iter.isNew {
134 return iter.First()
135 }
136 if iter.node == nil {
137 return false
138 }
139 // Reseek the previous key without allowing for an exact match if a force seek was requested. This results in the
140 // key greater than the previous one or an exhausted iterator if there is no such key.
141 if seekKey := iter.seekKey; seekKey != nil {
142 iter.seekKey = nil
143 return iter.seek(seekKey, false, true)
144 }
145 // When there is no right node walk the parents until the parent's right node is not equal to the previous child.
146 // This will be the next node.
147 if iter.node.right == nil {
148 parent := iter.parents.Pop()
149 for parent != nil && parent.right == iter.node {
150 iter.node = parent
151 parent = iter.parents.Pop()
152 }
153 iter.node = parent
154 return iter.limitIterator()
155 }
156 // There is a right node, so the next node is the left-most node down the right sub-tree.
157 iter.parents.Push(iter.node)
158 iter.node = iter.node.right
159 for node := iter.node.left; node != nil; node = node.left {
160 iter.parents.Push(iter.node)
161 iter.node = node
162 }
163 return iter.limitIterator()
164 }
165 166 // Prev moves the iterator to the previous key/value pair and returns false when the iterator is exhausted. When invoked
167 // on a newly created iterator it will position the iterator at the last item.
168 func (iter *Iterator) Prev() bool {
169 if iter.isNew {
170 return iter.Last()
171 }
172 if iter.node == nil {
173 return false
174 }
175 // Reseek the previous key without allowing for an exact match if a force seek was requested. This results in the
176 // key smaller than the previous one or an exhausted iterator if there is no such key.
177 if seekKey := iter.seekKey; seekKey != nil {
178 iter.seekKey = nil
179 return iter.seek(seekKey, false, false)
180 }
181 // When there is no left node walk the parents until the parent's left node is not equal to the previous child. This
182 // will be the previous node.
183 if iter.node.left == nil {
184 parent := iter.parents.Pop()
185 for parent != nil && parent.left == iter.node {
186 iter.node = parent
187 parent = iter.parents.Pop()
188 }
189 iter.node = parent
190 return iter.limitIterator()
191 }
192 // There is a left node, so the previous node is the right-most node down the left sub-tree.
193 iter.parents.Push(iter.node)
194 iter.node = iter.node.left
195 for node := iter.node.right; node != nil; node = node.right {
196 iter.parents.Push(iter.node)
197 iter.node = node
198 }
199 return iter.limitIterator()
200 }
201 202 // Seek moves the iterator to the first key/value pair with a key that is greater than or equal to the given key and
203 // returns true if successful.
204 func (iter *Iterator) Seek(key []byte) bool {
205 iter.isNew = false
206 return iter.seek(key, true, true)
207 }
208 209 // Key returns the key of the current key/value pair or nil when the iterator is exhausted. The caller should not modify
210 // the contents of the returned slice.
211 func (iter *Iterator) Key() []byte {
212 if iter.node == nil {
213 return nil
214 }
215 return iter.node.key
216 }
217 218 // Value returns the value of the current key/value pair or nil when the iterator is exhausted. The caller should not
219 // modify the contents of the returned slice.
220 func (iter *Iterator) Value() []byte {
221 if iter.node == nil {
222 return nil
223 }
224 return iter.node.value
225 }
226 227 // Valid indicates whether the iterator is positioned at a valid key/value pair. It will be considered invalid when the
228 // iterator is newly created or exhausted.
229 func (iter *Iterator) Valid() bool {
230 return iter.node != nil
231 }
232 233 // ForceReseek notifies the iterator that the underlying mutable treap has been updated, so the next call to Prev or
234 // Next needs to reseek in order to allow the iterator to continue working properly. NOTE: Calling this function when
235 // the iterator is associated with an immutable treap has no effect as you would expect.
236 func (iter *Iterator) ForceReseek() {
237 // Nothing to do when the iterator is associated with an immutable treap.
238 if iter.t == nil {
239 return
240 }
241 // Update the iterator root to the mutable treap root in case it changed.
242 iter.root = iter.t.root
243 // Set the seek key to the current node. This will force the Next/Prev functions to reseek, and thus properly
244 // reconstruct the iterator, on their next call.
245 if iter.node == nil {
246 iter.seekKey = nil
247 return
248 }
249 iter.seekKey = iter.node.key
250 }
251 252 // Iterator returns a new iterator for the mutable treap. The newly returned iterator is not pointing to a valid item
253 // until a call to one of the methods to position it is made. The start key and limit key parameters cause the iterator
254 // to be limited to a range of keys. The start key is inclusive and the limit key is exclusive. Either or both can be
255 // nil if the functionality is not desired.
256 //
257 // WARNING: The ForceSeek method must be called on the returned iterator if the treap is mutated. Failure to do so will
258 // cause the iterator to return unexpected keys and/or values.
259 //
260 // For example:
261 //
262 // iter := t.Iterator(nil, nil)
263 // for iter.Next() {
264 // if someCondition {
265 // t.Delete(iter.Key())
266 // iter.ForceReseek()
267 // }
268 // }
269 func (t *Mutable) Iterator(startKey, limitKey []byte) *Iterator {
270 iter := &Iterator{
271 t: t,
272 root: t.root,
273 isNew: true,
274 startKey: startKey,
275 limitKey: limitKey,
276 }
277 return iter
278 }
279 280 // Iterator returns a new iterator for the immutable treap. The newly returned iterator is not pointing to a valid item
281 // until a call to one of the methods to position it is made. The start key and limit key parameters cause the iterator
282 // to be limited to a range of keys.
283 //
284 // The start key is inclusive and the limit key is exclusive. Either or both can be nil if the functionality is not
285 // desired.
286 func (t *Immutable) Iterator(startKey, limitKey []byte) *Iterator {
287 iter := &Iterator{
288 root: t.root,
289 isNew: true,
290 startKey: startKey,
291 limitKey: limitKey,
292 }
293 return iter
294 }
295