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 }