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