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