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  }