github.com/MetalBlockchain/metalgo@v1.11.9/vms/platformvm/state/merged_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 "github.com/MetalBlockchain/metalgo/utils/heap"
     7  
     8  var _ StakerIterator = (*mergedIterator)(nil)
     9  
    10  type mergedIterator struct {
    11  	initialized bool
    12  	// heap only contains iterators that have been initialized and are not
    13  	// exhausted.
    14  	heap heap.Queue[StakerIterator]
    15  }
    16  
    17  // Returns an iterator that returns all of the elements of [stakers] in order.
    18  func NewMergedIterator(stakers ...StakerIterator) StakerIterator {
    19  	// Filter out iterators that are already exhausted.
    20  	i := 0
    21  	for i < len(stakers) {
    22  		staker := stakers[i]
    23  		if staker.Next() {
    24  			i++
    25  			continue
    26  		}
    27  		staker.Release()
    28  
    29  		newLength := len(stakers) - 1
    30  		stakers[i] = stakers[newLength]
    31  		stakers[newLength] = nil
    32  		stakers = stakers[:newLength]
    33  	}
    34  
    35  	it := &mergedIterator{
    36  		heap: heap.QueueOf(
    37  			func(a, b StakerIterator) bool {
    38  				return a.Value().Less(b.Value())
    39  			},
    40  			stakers...,
    41  		),
    42  	}
    43  
    44  	return it
    45  }
    46  
    47  func (it *mergedIterator) Next() bool {
    48  	if it.heap.Len() == 0 {
    49  		return false
    50  	}
    51  
    52  	if !it.initialized {
    53  		// Note that on the first call to Next() (i.e. here) we don't call
    54  		// Next() on the current iterator. This is because we already called
    55  		// Next() on each iterator in NewMergedIterator.
    56  		it.initialized = true
    57  		return true
    58  	}
    59  
    60  	// Update the heap root.
    61  	current, _ := it.heap.Peek()
    62  	if current.Next() {
    63  		// Calling Next() above modifies [current] so we fix the heap.
    64  		it.heap.Fix(0)
    65  		return true
    66  	}
    67  
    68  	// The old root is exhausted. Remove it from the heap.
    69  	current.Release()
    70  	it.heap.Pop()
    71  	return it.heap.Len() > 0
    72  }
    73  
    74  func (it *mergedIterator) Value() *Staker {
    75  	peek, _ := it.heap.Peek()
    76  	return peek.Value()
    77  }
    78  
    79  func (it *mergedIterator) Release() {
    80  	for it.heap.Len() > 0 {
    81  		removed, _ := it.heap.Pop()
    82  		removed.Release()
    83  	}
    84  }
    85  
    86  func (it *mergedIterator) Len() int {
    87  	return it.heap.Len()
    88  }