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 }