github.com/MetalBlockchain/metalgo@v1.11.9/snow/consensus/snowman/bootstrapper/sampler.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package bootstrapper
     5  
     6  import (
     7  	"errors"
     8  
     9  	"github.com/MetalBlockchain/metalgo/utils/math"
    10  	"github.com/MetalBlockchain/metalgo/utils/sampler"
    11  	"github.com/MetalBlockchain/metalgo/utils/set"
    12  )
    13  
    14  var errUnexpectedSamplerFailure = errors.New("unexpected sampler failure")
    15  
    16  // Sample keys from [elements] uniformly by weight without replacement. The
    17  // returned set will have size less than or equal to [maxSize]. This function
    18  // will error if the sum of all weights overflows.
    19  func Sample[T comparable](elements map[T]uint64, maxSize int) (set.Set[T], error) {
    20  	var (
    21  		keys        = make([]T, len(elements))
    22  		weights     = make([]uint64, len(elements))
    23  		totalWeight uint64
    24  		err         error
    25  	)
    26  	i := 0
    27  	for key, weight := range elements {
    28  		keys[i] = key
    29  		weights[i] = weight
    30  		totalWeight, err = math.Add64(totalWeight, weight)
    31  		if err != nil {
    32  			return nil, err
    33  		}
    34  		i++
    35  	}
    36  
    37  	sampler := sampler.NewWeightedWithoutReplacement()
    38  	if err := sampler.Initialize(weights); err != nil {
    39  		return nil, err
    40  	}
    41  
    42  	maxSize = int(min(uint64(maxSize), totalWeight))
    43  	indices, ok := sampler.Sample(maxSize)
    44  	if !ok {
    45  		return nil, errUnexpectedSamplerFailure
    46  	}
    47  
    48  	sampledElements := set.NewSet[T](maxSize)
    49  	for _, index := range indices {
    50  		sampledElements.Add(keys[index])
    51  	}
    52  	return sampledElements, nil
    53  }