github.com/ava-labs/avalanchego@v1.11.11/vms/platformvm/state/staker_diff_iterator.go (about) 1 // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved. 2 // See the file LICENSE for licensing terms. 3 4 package state 5 6 import ( 7 "github.com/ava-labs/avalanchego/utils/heap" 8 "github.com/ava-labs/avalanchego/utils/iterator" 9 "github.com/ava-labs/avalanchego/vms/platformvm/txs" 10 ) 11 12 var ( 13 _ StakerDiffIterator = (*stakerDiffIterator)(nil) 14 _ iterator.Iterator[*Staker] = (*mutableStakerIterator)(nil) 15 ) 16 17 // StakerDiffIterator is an iterator that iterates over the events that will be 18 // performed on the current staker set. 19 // 20 // There are two event types affecting current staker set, removal of an 21 // existing staker and addition of a new staker from the pending set. 22 // 23 // The ordering of operations is: 24 // - Staker operations are performed in order of their [NextTime]. 25 // - If operations have the same [NextTime], stakers are first added to the 26 // current staker set, then removed. 27 // - Further ties are broken by *Staker.Less(), returning the lesser staker 28 // first. 29 type StakerDiffIterator interface { 30 Next() bool 31 // Returns: 32 // - The staker that is changing 33 // - True if the staker is being added to the current staker set, false if 34 // the staker is being removed from the current staker set 35 Value() (*Staker, bool) 36 Release() 37 } 38 39 type stakerDiffIterator struct { 40 currentIteratorExhausted bool 41 currentIterator *mutableStakerIterator 42 43 pendingIteratorExhausted bool 44 pendingIterator iterator.Iterator[*Staker] 45 46 modifiedStaker *Staker 47 isAdded bool 48 } 49 50 func NewStakerDiffIterator(currentIterator, pendingIterator iterator.Iterator[*Staker]) StakerDiffIterator { 51 mutableCurrentIterator := newMutableStakerIterator(currentIterator) 52 return &stakerDiffIterator{ 53 currentIteratorExhausted: !mutableCurrentIterator.Next(), 54 currentIterator: mutableCurrentIterator, 55 pendingIteratorExhausted: !pendingIterator.Next(), 56 pendingIterator: pendingIterator, 57 } 58 } 59 60 func (it *stakerDiffIterator) Next() bool { 61 switch { 62 case it.currentIteratorExhausted && it.pendingIteratorExhausted: 63 return false 64 case it.currentIteratorExhausted: 65 it.advancePending() 66 case it.pendingIteratorExhausted: 67 it.advanceCurrent() 68 default: 69 nextStakerRemoved := it.currentIterator.Value() 70 nextStakerAdded := it.pendingIterator.Value() 71 // If the next operations share the same time, we default to adding the 72 // staker to the current staker set. This means that we default to 73 // advancing the pending iterator. 74 if nextStakerRemoved.EndTime.Before(nextStakerAdded.StartTime) { 75 it.advanceCurrent() 76 } else { 77 it.advancePending() 78 } 79 } 80 return true 81 } 82 83 func (it *stakerDiffIterator) Value() (*Staker, bool) { 84 return it.modifiedStaker, it.isAdded 85 } 86 87 func (it *stakerDiffIterator) Release() { 88 it.currentIteratorExhausted = true 89 it.currentIterator.Release() 90 it.pendingIteratorExhausted = true 91 it.pendingIterator.Release() 92 it.modifiedStaker = nil 93 } 94 95 func (it *stakerDiffIterator) advanceCurrent() { 96 it.modifiedStaker = it.currentIterator.Value() 97 it.isAdded = false 98 it.currentIteratorExhausted = !it.currentIterator.Next() 99 } 100 101 func (it *stakerDiffIterator) advancePending() { 102 it.modifiedStaker = it.pendingIterator.Value() 103 it.isAdded = true 104 it.pendingIteratorExhausted = !it.pendingIterator.Next() 105 106 toRemove := *it.modifiedStaker 107 toRemove.NextTime = toRemove.EndTime 108 toRemove.Priority = txs.PendingToCurrentPriorities[toRemove.Priority] 109 it.currentIteratorExhausted = false 110 it.currentIterator.Add(&toRemove) 111 } 112 113 type mutableStakerIterator struct { 114 iteratorExhausted bool 115 iterator iterator.Iterator[*Staker] 116 heap heap.Queue[*Staker] 117 } 118 119 func newMutableStakerIterator(iterator iterator.Iterator[*Staker]) *mutableStakerIterator { 120 return &mutableStakerIterator{ 121 iteratorExhausted: !iterator.Next(), 122 iterator: iterator, 123 heap: heap.NewQueue((*Staker).Less), 124 } 125 } 126 127 // Add should not be called until after Next has been called at least once. 128 func (it *mutableStakerIterator) Add(staker *Staker) { 129 it.heap.Push(staker) 130 } 131 132 func (it *mutableStakerIterator) Next() bool { 133 // The only time the heap should be empty - is when the iterator is 134 // exhausted or uninitialized. 135 if it.heap.Len() > 0 { 136 it.heap.Pop() 137 } 138 139 // If the iterator is exhausted, the only elements left to iterate over are 140 // in the heap. 141 if it.iteratorExhausted { 142 return it.heap.Len() > 0 143 } 144 145 // If the heap doesn't contain the next staker to return, we need to move 146 // the next element from the iterator into the heap. 147 nextIteratorStaker := it.iterator.Value() 148 peek, ok := it.heap.Peek() 149 if !ok || nextIteratorStaker.Less(peek) { 150 it.Add(nextIteratorStaker) 151 it.iteratorExhausted = !it.iterator.Next() 152 } 153 return true 154 } 155 156 func (it *mutableStakerIterator) Value() *Staker { 157 peek, _ := it.heap.Peek() 158 return peek 159 } 160 161 func (it *mutableStakerIterator) Release() { 162 it.iteratorExhausted = true 163 it.iterator.Release() 164 it.heap = heap.NewQueue((*Staker).Less) 165 }