github.com/songzhibin97/go-baseutils@v0.0.2-0.20240302024150-487d8ce9c082/structure/trees/avltree/iterator.go (about)

     1  package avltree
     2  
     3  import "github.com/songzhibin97/go-baseutils/structure/containers"
     4  
     5  // Assert Iterator implementation
     6  var _ containers.ReverseIteratorWithKey[any, any] = (*Iterator[any, any])(nil)
     7  
     8  // Iterator holding the iterator's state
     9  type Iterator[K any, V any] struct {
    10  	tree     *Tree[K, V]
    11  	node     *Node[K, V]
    12  	position position
    13  	zeroK    K
    14  	zeroV    V
    15  }
    16  
    17  type position byte
    18  
    19  const (
    20  	begin, between, end position = 0, 1, 2
    21  )
    22  
    23  // Iterator returns a stateful iterator whose elements are key/value pairs.
    24  func (tree *Tree[K, V]) Iterator() containers.ReverseIteratorWithKey[K, V] {
    25  	return &Iterator[K, V]{tree: tree, node: nil, position: begin}
    26  }
    27  
    28  // Next moves the iterator to the next element and returns true if there was a next element in the container.
    29  // If Next() returns true, then next element's key and value can be retrieved by Key() and Value().
    30  // If Next() was called for the first time, then it will point the iterator to the first element if it exists.
    31  // Modifies the state of the iterator.
    32  func (iterator *Iterator[K, V]) Next() bool {
    33  	switch iterator.position {
    34  	case begin:
    35  		iterator.position = between
    36  		iterator.node = iterator.tree.Left()
    37  	case between:
    38  		iterator.node = iterator.node.Next()
    39  	}
    40  
    41  	if iterator.node == nil {
    42  		iterator.position = end
    43  		return false
    44  	}
    45  	return true
    46  }
    47  
    48  // Prev moves the iterator to the next element and returns true if there was a previous element in the container.
    49  // If Prev() returns true, then next element's key and value can be retrieved by Key() and Value().
    50  // If Prev() was called for the first time, then it will point the iterator to the first element if it exists.
    51  // Modifies the state of the iterator.
    52  func (iterator *Iterator[K, V]) Prev() bool {
    53  	switch iterator.position {
    54  	case end:
    55  		iterator.position = between
    56  		iterator.node = iterator.tree.Right()
    57  	case between:
    58  		iterator.node = iterator.node.Prev()
    59  	}
    60  
    61  	if iterator.node == nil {
    62  		iterator.position = begin
    63  		return false
    64  	}
    65  	return true
    66  }
    67  
    68  // Value returns the current element's value.
    69  // Does not modify the state of the iterator.
    70  func (iterator *Iterator[K, V]) Value() V {
    71  	if iterator.node == nil {
    72  		return iterator.zeroV
    73  	}
    74  	return iterator.node.Value
    75  }
    76  
    77  // Key returns the current element's key.
    78  // Does not modify the state of the iterator.
    79  func (iterator *Iterator[K, V]) Key() K {
    80  	if iterator.node == nil {
    81  		return iterator.zeroK
    82  	}
    83  	return iterator.node.Key
    84  }
    85  
    86  // Node returns the current element's node.
    87  // Does not modify the state of the iterator.
    88  func (iterator *Iterator[K, V]) Node() *Node[K, V] {
    89  	return iterator.node
    90  }
    91  
    92  // Begin resets the iterator to its initial state (one-before-first)
    93  // Call Next() to fetch the first element if any.
    94  func (iterator *Iterator[K, V]) Begin() {
    95  	iterator.node = nil
    96  	iterator.position = begin
    97  }
    98  
    99  // End moves the iterator past the last element (one-past-the-end).
   100  // Call Prev() to fetch the last element if any.
   101  func (iterator *Iterator[K, V]) End() {
   102  	iterator.node = nil
   103  	iterator.position = end
   104  }
   105  
   106  // First moves the iterator to the first element and returns true if there was a first element in the container.
   107  // If First() returns true, then first element's key and value can be retrieved by Key() and Value().
   108  // Modifies the state of the iterator
   109  func (iterator *Iterator[K, V]) First() bool {
   110  	iterator.Begin()
   111  	return iterator.Next()
   112  }
   113  
   114  // Last moves the iterator to the last element and returns true if there was a last element in the container.
   115  // If Last() returns true, then last element's key and value can be retrieved by Key() and Value().
   116  // Modifies the state of the iterator.
   117  func (iterator *Iterator[K, V]) Last() bool {
   118  	iterator.End()
   119  	return iterator.Prev()
   120  }
   121  
   122  // NextTo moves the iterator to the next element from current position that satisfies the condition given by the
   123  // passed function, and returns true if there was a next element in the container.
   124  // If NextTo() returns true, then next element's key and value can be retrieved by Key() and Value().
   125  // Modifies the state of the iterator.
   126  func (iterator *Iterator[K, V]) NextTo(f func(key K, value V) bool) bool {
   127  	for iterator.Next() {
   128  		key, value := iterator.Key(), iterator.Value()
   129  		if f(key, value) {
   130  			return true
   131  		}
   132  	}
   133  	return false
   134  }
   135  
   136  // PrevTo moves the iterator to the previous element from current position that satisfies the condition given by the
   137  // passed function, and returns true if there was a next element in the container.
   138  // If PrevTo() returns true, then next element's key and value can be retrieved by Key() and Value().
   139  // Modifies the state of the iterator.
   140  func (iterator *Iterator[K, V]) PrevTo(f func(key K, value V) bool) bool {
   141  	for iterator.Prev() {
   142  		key, value := iterator.Key(), iterator.Value()
   143  		if f(key, value) {
   144  			return true
   145  		}
   146  	}
   147  	return false
   148  }