github.com/MetalBlockchain/metalgo@v1.11.9/vms/platformvm/state/tree_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  	"sync"
     8  
     9  	"github.com/google/btree"
    10  )
    11  
    12  const defaultTreeDegree = 2
    13  
    14  var _ StakerIterator = (*treeIterator)(nil)
    15  
    16  type treeIterator struct {
    17  	current     *Staker
    18  	next        chan *Staker
    19  	releaseOnce sync.Once
    20  	release     chan struct{}
    21  	wg          sync.WaitGroup
    22  }
    23  
    24  // NewTreeIterator returns a new iterator of the stakers in [tree] in ascending
    25  // order. Note that it isn't safe to modify [tree] while iterating over it.
    26  func NewTreeIterator(tree *btree.BTreeG[*Staker]) StakerIterator {
    27  	if tree == nil {
    28  		return EmptyIterator
    29  	}
    30  	it := &treeIterator{
    31  		next:    make(chan *Staker),
    32  		release: make(chan struct{}),
    33  	}
    34  	it.wg.Add(1)
    35  	go func() {
    36  		defer it.wg.Done()
    37  		tree.Ascend(func(i *Staker) bool {
    38  			select {
    39  			case it.next <- i:
    40  				return true
    41  			case <-it.release:
    42  				return false
    43  			}
    44  		})
    45  		close(it.next)
    46  	}()
    47  	return it
    48  }
    49  
    50  func (i *treeIterator) Next() bool {
    51  	next, ok := <-i.next
    52  	i.current = next
    53  	return ok
    54  }
    55  
    56  func (i *treeIterator) Value() *Staker {
    57  	return i.current
    58  }
    59  
    60  func (i *treeIterator) Release() {
    61  	i.releaseOnce.Do(func() {
    62  		close(i.release)
    63  	})
    64  	i.wg.Wait()
    65  }