treapiter.go

// Copyright (c) 2015-2016 The btcsuite developers
// Copyright (c) 2016-2022 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"

// 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.
	if 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
}