Simplify sampler interface (#3026)

StephenButtolph · web-flow · commit ddd6d25379e2 · 2024-05-17T23:14:49.000Z
diff --git a/network/ip_tracker.go b/network/ip_tracker.go
@@ -400,8 +400,8 @@ func (i *ipTracker) GetGossipableIPs(
 
 	uniform.Initialize(uint64(len(i.gossipableIPs)))
 	for len(ips) < maxNumIPs {
-		index, err := uniform.Next()
-		if err != nil {
+		index, hasNext := uniform.Next()
+		if !hasNext {
 			return ips
 		}
 
diff --git a/network/p2p/validators.go b/network/p2p/validators.go
@@ -125,8 +125,8 @@ func (v *Validators) Sample(ctx context.Context, limit int) []ids.NodeID {
 
 	uniform.Initialize(uint64(len(v.validatorList)))
 	for len(sampled) < limit {
-		i, err := uniform.Next()
-		if err != nil {
+		i, hasNext := uniform.Next()
+		if !hasNext {
 			break
 		}
 
diff --git a/network/peer/set.go b/network/peer/set.go
@@ -124,8 +124,8 @@ func (s *peerSet) Sample(n int, precondition func(Peer) bool) []Peer {
 
 	peers := make([]Peer, 0, n)
 	for len(peers) < n {
-		index, err := sampler.Next()
-		if err != nil {
+		index, hasNext := sampler.Next()
+		if !hasNext {
 			// We have run out of peers to attempt to sample.
 			break
 		}
diff --git a/snow/consensus/snowman/bootstrapper/sampler.go b/snow/consensus/snowman/bootstrapper/sampler.go
@@ -4,11 +4,15 @@
 package bootstrapper
 
 import (
+	"errors"
+
 	"github.com/ava-labs/avalanchego/utils/math"
 	"github.com/ava-labs/avalanchego/utils/sampler"
 	"github.com/ava-labs/avalanchego/utils/set"
 )
 
+var errUnexpectedSamplerFailure = errors.New("unexpected sampler failure")
+
 // Sample keys from [elements] uniformly by weight without replacement. The
 // returned set will have size less than or equal to [maxSize]. This function
 // will error if the sum of all weights overflows.
@@ -36,9 +40,9 @@ func Sample[T comparable](elements map[T]uint64, maxSize int) (set.Set[T], error
 	}
 
 	maxSize = int(min(uint64(maxSize), totalWeight))
-	indices, err := sampler.Sample(maxSize)
-	if err != nil {
-		return nil, err
+	indices, ok := sampler.Sample(maxSize)
+	if !ok {
+		return nil, errUnexpectedSamplerFailure
 	}
 
 	sampledElements := set.NewSet[T](maxSize)
diff --git a/snow/validators/manager_test.go b/snow/validators/manager_test.go
@@ -11,7 +11,6 @@ import (
 
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/utils/crypto/bls"
-	"github.com/ava-labs/avalanchego/utils/sampler"
 	"github.com/ava-labs/avalanchego/utils/set"
 
 	safemath "github.com/ava-labs/avalanchego/utils/math"
@@ -396,7 +395,7 @@ func TestSample(t *testing.T) {
 	require.Equal([]ids.NodeID{nodeID0}, sampled)
 
 	_, err = m.Sample(subnetID, 2)
-	require.ErrorIs(err, sampler.ErrOutOfRange)
+	require.ErrorIs(err, errInsufficientWeight)
 
 	nodeID1 := ids.GenerateTestNodeID()
 	require.NoError(m.AddStaker(subnetID, nodeID1, nil, ids.Empty, math.MaxInt64-1))
diff --git a/snow/validators/set.go b/snow/validators/set.go
@@ -23,6 +23,7 @@ var (
 	errDuplicateValidator   = errors.New("duplicate validator")
 	errMissingValidator     = errors.New("missing validator")
 	errTotalWeightNotUint64 = errors.New("total weight is not a uint64")
+	errInsufficientWeight   = errors.New("insufficient weight")
 )
 
 // newSet returns a new, empty set of validators.
@@ -257,9 +258,9 @@ func (s *vdrSet) sample(size int) ([]ids.NodeID, error) {
 		s.samplerInitialized = true
 	}
 
-	indices, err := s.sampler.Sample(size)
-	if err != nil {
-		return nil, err
+	indices, ok := s.sampler.Sample(size)
+	if !ok {
+		return nil, errInsufficientWeight
 	}
 
 	list := make([]ids.NodeID, size)
diff --git a/snow/validators/set_test.go b/snow/validators/set_test.go
@@ -11,7 +11,6 @@ import (
 
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/utils/crypto/bls"
-	"github.com/ava-labs/avalanchego/utils/sampler"
 	"github.com/ava-labs/avalanchego/utils/set"
 
 	safemath "github.com/ava-labs/avalanchego/utils/math"
@@ -343,7 +342,7 @@ func TestSetSample(t *testing.T) {
 	require.Equal([]ids.NodeID{nodeID0}, sampled)
 
 	_, err = s.Sample(2)
-	require.ErrorIs(err, sampler.ErrOutOfRange)
+	require.ErrorIs(err, errInsufficientWeight)
 
 	nodeID1 := ids.GenerateTestNodeID()
 	require.NoError(s.Add(nodeID1, nil, ids.Empty, math.MaxInt64-1))
diff --git a/utils/sampler/uniform.go b/utils/sampler/uniform.go
@@ -7,12 +7,12 @@ package sampler
 type Uniform interface {
 	Initialize(sampleRange uint64)
 	// Sample returns length numbers in the range [0,sampleRange). If there
-	// aren't enough numbers in the range, an error is returned. If length is
+	// aren't enough numbers in the range, false is returned. If length is
 	// negative the implementation may panic.
-	Sample(length int) ([]uint64, error)
+	Sample(length int) ([]uint64, bool)
 
+	Next() (uint64, bool)
 	Reset()
-	Next() (uint64, error)
 }
 
 // NewUniform returns a new sampler
diff --git a/utils/sampler/uniform_best.go b/utils/sampler/uniform_best.go
@@ -56,7 +56,7 @@ samplerLoop:
 
 		start := s.clock.Time()
 		for i := 0; i < s.benchmarkIterations; i++ {
-			if _, err := sampler.Sample(sampleSize); err != nil {
+			if _, ok := sampler.Sample(sampleSize); !ok {
 				continue samplerLoop
 			}
 		}
diff --git a/utils/sampler/uniform_replacer.go b/utils/sampler/uniform_replacer.go
@@ -36,34 +36,34 @@ func (s *uniformReplacer) Initialize(length uint64) {
 	s.drawsCount = 0
 }
 
-func (s *uniformReplacer) Sample(count int) ([]uint64, error) {
+func (s *uniformReplacer) Sample(count int) ([]uint64, bool) {
 	s.Reset()
 
 	results := make([]uint64, count)
 	for i := 0; i < count; i++ {
-		ret, err := s.Next()
-		if err != nil {
-			return nil, err
+		ret, hasNext := s.Next()
+		if !hasNext {
+			return nil, false
 		}
 		results[i] = ret
 	}
-	return results, nil
+	return results, true
 }
 
 func (s *uniformReplacer) Reset() {
 	clear(s.drawn)
 	s.drawsCount = 0
 }
 
-func (s *uniformReplacer) Next() (uint64, error) {
+func (s *uniformReplacer) Next() (uint64, bool) {
 	if s.drawsCount >= s.length {
-		return 0, ErrOutOfRange
+		return 0, false
 	}
 
 	draw := s.rng.Uint64Inclusive(s.length-1-s.drawsCount) + s.drawsCount
 	ret := s.drawn.get(draw, draw)
 	s.drawn[draw] = s.drawn.get(s.drawsCount, s.drawsCount)
 	s.drawsCount++
 
-	return ret, nil
+	return ret, true
 }
diff --git a/utils/sampler/uniform_resample.go b/utils/sampler/uniform_resample.go
@@ -23,28 +23,28 @@ func (s *uniformResample) Initialize(length uint64) {
 	s.drawn = make(map[uint64]struct{})
 }
 
-func (s *uniformResample) Sample(count int) ([]uint64, error) {
+func (s *uniformResample) Sample(count int) ([]uint64, bool) {
 	s.Reset()
 
 	results := make([]uint64, count)
 	for i := 0; i < count; i++ {
-		ret, err := s.Next()
-		if err != nil {
-			return nil, err
+		ret, hasNext := s.Next()
+		if !hasNext {
+			return nil, false
 		}
 		results[i] = ret
 	}
-	return results, nil
+	return results, true
 }
 
 func (s *uniformResample) Reset() {
 	clear(s.drawn)
 }
 
-func (s *uniformResample) Next() (uint64, error) {
+func (s *uniformResample) Next() (uint64, bool) {
 	i := uint64(len(s.drawn))
 	if i >= s.length {
-		return 0, ErrOutOfRange
+		return 0, false
 	}
 
 	for {
@@ -53,6 +53,6 @@ func (s *uniformResample) Next() (uint64, error) {
 			continue
 		}
 		s.drawn[draw] = struct{}{}
-		return draw, nil
+		return draw, true
 	}
 }
diff --git a/utils/sampler/uniform_test.go b/utils/sampler/uniform_test.go
@@ -83,8 +83,8 @@ func UniformInitializeMaxUint64Test(t *testing.T, s Uniform) {
 	s.Initialize(math.MaxUint64)
 
 	for {
-		val, err := s.Next()
-		require.NoError(t, err)
+		val, hasNext := s.Next()
+		require.True(t, hasNext)
 
 		if val > math.MaxInt64 {
 			break
@@ -95,17 +95,17 @@ func UniformInitializeMaxUint64Test(t *testing.T, s Uniform) {
 func UniformOutOfRangeTest(t *testing.T, s Uniform) {
 	s.Initialize(0)
 
-	_, err := s.Sample(1)
-	require.ErrorIs(t, err, ErrOutOfRange)
+	_, ok := s.Sample(1)
+	require.False(t, ok)
 }
 
 func UniformEmptyTest(t *testing.T, s Uniform) {
 	require := require.New(t)
 
 	s.Initialize(1)
 
-	val, err := s.Sample(0)
-	require.NoError(err)
+	val, ok := s.Sample(0)
+	require.True(ok)
 	require.Empty(val)
 }
 
@@ -114,8 +114,8 @@ func UniformSingletonTest(t *testing.T, s Uniform) {
 
 	s.Initialize(1)
 
-	val, err := s.Sample(1)
-	require.NoError(err)
+	val, ok := s.Sample(1)
+	require.True(ok)
 	require.Equal([]uint64{0}, val)
 }
 
@@ -124,8 +124,8 @@ func UniformDistributionTest(t *testing.T, s Uniform) {
 
 	s.Initialize(3)
 
-	val, err := s.Sample(3)
-	require.NoError(err)
+	val, ok := s.Sample(3)
+	require.True(ok)
 
 	slices.Sort(val)
 	require.Equal([]uint64{0, 1, 2}, val)
@@ -134,8 +134,8 @@ func UniformDistributionTest(t *testing.T, s Uniform) {
 func UniformOverSampleTest(t *testing.T, s Uniform) {
 	s.Initialize(3)
 
-	_, err := s.Sample(4)
-	require.ErrorIs(t, err, ErrOutOfRange)
+	_, ok := s.Sample(4)
+	require.False(t, ok)
 }
 
 func UniformLazilySample(t *testing.T, s Uniform) {
@@ -146,15 +146,15 @@ func UniformLazilySample(t *testing.T, s Uniform) {
 	for j := 0; j < 2; j++ {
 		sampled := map[uint64]bool{}
 		for i := 0; i < 3; i++ {
-			val, err := s.Next()
-			require.NoError(err)
+			val, hasNext := s.Next()
+			require.True(hasNext)
 			require.False(sampled[val])
 
 			sampled[val] = true
 		}
 
-		_, err := s.Next()
-		require.ErrorIs(err, ErrOutOfRange)
+		_, hasNext := s.Next()
+		require.False(hasNext)
 
 		s.Reset()
 	}
diff --git a/utils/sampler/weighted.go b/utils/sampler/weighted.go
@@ -3,15 +3,11 @@
 
 package sampler
 
-import "errors"
-
-var ErrOutOfRange = errors.New("out of range")
-
 // Weighted defines how to sample a specified valued based on a provided
 // weighted distribution
 type Weighted interface {
 	Initialize(weights []uint64) error
-	Sample(sampleValue uint64) (int, error)
+	Sample(sampleValue uint64) (int, bool)
 }
 
 // NewWeighted returns a new sampler
diff --git a/utils/sampler/weighted_array.go b/utils/sampler/weighted_array.go
@@ -81,9 +81,9 @@ func (s *weightedArray) Initialize(weights []uint64) error {
 	return nil
 }
 
-func (s *weightedArray) Sample(value uint64) (int, error) {
+func (s *weightedArray) Sample(value uint64) (int, bool) {
 	if len(s.arr) == 0 || s.arr[len(s.arr)-1].cumulativeWeight <= value {
-		return 0, ErrOutOfRange
+		return 0, false
 	}
 	minIndex := 0
 	maxIndex := len(s.arr) - 1
@@ -98,7 +98,7 @@ func (s *weightedArray) Sample(value uint64) (int, error) {
 		currentElem := s.arr[index]
 		currentWeight := currentElem.cumulativeWeight
 		if previousWeight <= value && value < currentWeight {
-			return currentElem.index, nil
+			return currentElem.index, true
 		}
 
 		if value < previousWeight {
diff --git a/utils/sampler/weighted_best.go b/utils/sampler/weighted_best.go
@@ -60,7 +60,7 @@ samplerLoop:
 
 		start := s.clock.Time()
 		for _, sample := range samples {
-			if _, err := sampler.Sample(sample); err != nil {
+			if _, ok := sampler.Sample(sample); !ok {
 				continue samplerLoop
 			}
 		}
diff --git a/utils/sampler/weighted_heap.go b/utils/sampler/weighted_heap.go
@@ -80,17 +80,17 @@ func (s *weightedHeap) Initialize(weights []uint64) error {
 	return nil
 }
 
-func (s *weightedHeap) Sample(value uint64) (int, error) {
+func (s *weightedHeap) Sample(value uint64) (int, bool) {
 	if len(s.heap) == 0 || s.heap[0].cumulativeWeight <= value {
-		return 0, ErrOutOfRange
+		return 0, false
 	}
 
 	index := 0
 	for {
 		currentElement := s.heap[index]
 		currentWeight := currentElement.weight
 		if value < currentWeight {
-			return currentElement.index, nil
+			return currentElement.index, true
 		}
 		value -= currentWeight
 
diff --git a/utils/sampler/weighted_linear.go b/utils/sampler/weighted_linear.go
diff --git a/utils/sampler/weighted_test.go b/utils/sampler/weighted_test.go
diff --git a/utils/sampler/weighted_uniform.go b/utils/sampler/weighted_uniform.go
diff --git a/utils/sampler/weighted_without_replacement.go b/utils/sampler/weighted_without_replacement.go
diff --git a/utils/sampler/weighted_without_replacement_generic.go b/utils/sampler/weighted_without_replacement_generic.go
diff --git a/utils/sampler/weighted_without_replacement_test.go b/utils/sampler/weighted_without_replacement_test.go
diff --git a/vms/avm/environment_test.go b/vms/avm/environment_test.go
diff --git a/vms/proposervm/proposer/windower.go b/vms/proposervm/proposer/windower.go

Original file line number	Diff line number	Diff line change
`@@ -400,8 +400,8 @@ func (i *ipTracker) GetGossipableIPs(`
`400`	`400`
`401`	`401`	`uniform.Initialize(uint64(len(i.gossipableIPs)))`
`402`	`402`	`for len(ips) < maxNumIPs {`
`403`		`- index, err := uniform.Next()`
`404`		`- if err != nil {`
	`403`	`+ index, hasNext := uniform.Next()`
	`404`	`+ if !hasNext {`
`405`	`405`	`return ips`
`406`	`406`	`}`
`407`	`407`
Original file line number	Diff line number	Diff line change
`@@ -125,8 +125,8 @@ func (v *Validators) Sample(ctx context.Context, limit int) []ids.NodeID {`
`125`	`125`
`126`	`126`	`uniform.Initialize(uint64(len(v.validatorList)))`
`127`	`127`	`for len(sampled) < limit {`
`128`		`- i, err := uniform.Next()`
`129`		`- if err != nil {`
	`128`	`+ i, hasNext := uniform.Next()`
	`129`	`+ if !hasNext {`
`130`	`130`	`break`
`131`	`131`	`}`
`132`	`132`
Original file line number	Diff line number	Diff line change
`@@ -23,6 +23,7 @@ var (`
`23`	`23`	`errDuplicateValidator = errors.New("duplicate validator")`
`24`	`24`	`errMissingValidator = errors.New("missing validator")`
`25`	`25`	`errTotalWeightNotUint64 = errors.New("total weight is not a uint64")`
	`26`	`+ errInsufficientWeight = errors.New("insufficient weight")`
`26`	`27`	`)`
`27`	`28`
`28`	`29`	`// newSet returns a new, empty set of validators.`
`@@ -257,9 +258,9 @@ func (s *vdrSet) sample(size int) ([]ids.NodeID, error) {`
`257`	`258`	`s.samplerInitialized = true`
`258`	`259`	`}`
`259`	`260`
`260`		`- indices, err := s.sampler.Sample(size)`
`261`		`- if err != nil {`
`262`		`- return nil, err`
	`261`	`+ indices, ok := s.sampler.Sample(size)`
	`262`	`+ if !ok {`
	`263`	`+ return nil, errInsufficientWeight`
`263`	`264`	`}`
`264`	`265`
`265`	`266`	`list := make([]ids.NodeID, size)`
Original file line number	Diff line number	Diff line change
`@@ -56,7 +56,7 @@ samplerLoop:`
`56`	`56`
`57`	`57`	`start := s.clock.Time()`
`58`	`58`	`for i := 0; i < s.benchmarkIterations; i++ {`
`59`		`- if _, err := sampler.Sample(sampleSize); err != nil {`
	`59`	`+ if _, ok := sampler.Sample(sampleSize); !ok {`
`60`	`60`	`continue samplerLoop`
`61`	`61`	`}`
`62`	`62`	`}`
Original file line number	Diff line number	Diff line change
`@@ -36,34 +36,34 @@ func (s *uniformReplacer) Initialize(length uint64) {`
`36`	`36`	`s.drawsCount = 0`
`37`	`37`	`}`
`38`	`38`
`39`		`-func (s *uniformReplacer) Sample(count int) ([]uint64, error) {`
	`39`	`+func (s *uniformReplacer) Sample(count int) ([]uint64, bool) {`
`40`	`40`	`s.Reset()`
`41`	`41`
`42`	`42`	`results := make([]uint64, count)`
`43`	`43`	`for i := 0; i < count; i++ {`
`44`		`- ret, err := s.Next()`
`45`		`- if err != nil {`
`46`		`- return nil, err`
	`44`	`+ ret, hasNext := s.Next()`
	`45`	`+ if !hasNext {`
	`46`	`+ return nil, false`
`47`	`47`	`}`
`48`	`48`	`results[i] = ret`
`49`	`49`	`}`
`50`		`- return results, nil`
	`50`	`+ return results, true`
`51`	`51`	`}`
`52`	`52`
`53`	`53`	`func (s *uniformReplacer) Reset() {`
`54`	`54`	`clear(s.drawn)`
`55`	`55`	`s.drawsCount = 0`
`56`	`56`	`}`
`57`	`57`
`58`		`-func (s *uniformReplacer) Next() (uint64, error) {`
	`58`	`+func (s *uniformReplacer) Next() (uint64, bool) {`
`59`	`59`	`if s.drawsCount >= s.length {`
`60`		`- return 0, ErrOutOfRange`
	`60`	`+ return 0, false`
`61`	`61`	`}`
`62`	`62`
`63`	`63`	`draw := s.rng.Uint64Inclusive(s.length-1-s.drawsCount) + s.drawsCount`
`64`	`64`	`ret := s.drawn.get(draw, draw)`
`65`	`65`	`s.drawn[draw] = s.drawn.get(s.drawsCount, s.drawsCount)`
`66`	`66`	`s.drawsCount++`
`67`	`67`
`68`		`- return ret, nil`
	`68`	`+ return ret, true`
`69`	`69`	`}`