vendor/github.com/btcsuite/btcd/database/internal/treap/treapiter.go

   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
   5 package treap
   6
   7 import "bytes"
   8
   9 // Iterator represents an iterator for forwards and backwards iteration over
  10 // the contents of a treap (mutable or immutable).
  11 type 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.
  25 func (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.
  55 func (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.
 116 func (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.
 140 func (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.
 164 func (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.
 207 func (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.
 250 func (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.
 258 func (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.
 268 func (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.
 277 func (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.
 287 func (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 //   }
 328 func (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.
 346 func (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 }