ceremonyclient/pebble/internal/randvar/deck.go

// Copyright 2019 The LevelDB-Go and Pebble Authors. All rights reserved. Use
// of this source code is governed by a BSD-style license that can be found in
// the LICENSE file.

package randvar

import (
	"sync"

	"golang.org/x/exp/rand"
)

// Deck is a random number generator that generates numbers in the range
// [0,len(weights)-1] where the probability of i is
// weights(i)/sum(weights). Unlike Weighted, the weights are specified as
// integers and used in a deck-of-cards style random number selection which
// ensures that each element is returned with a desired frequency within the
// size of the deck.
type Deck struct {
	rng *rand.Rand
	mu  struct {
		sync.Mutex
		index int
		deck  []int
	}
}

// NewDeck returns a new deck random number generator.
func NewDeck(rng *rand.Rand, weights ...int) *Deck {
	var sum int
	for i := range weights {
		sum += weights[i]
	}
	deck := make([]int, 0, sum)
	for i := range weights {
		for j := 0; j < weights[i]; j++ {
			deck = append(deck, i)
		}
	}
	d := &Deck{
		rng: ensureRand(rng),
	}
	d.mu.index = len(deck)
	d.mu.deck = deck
	return d
}

// Int returns a random number in the range [0,len(weights)-1] where the
// probability of i is weights(i)/sum(weights).
func (d *Deck) Int() int {
	d.mu.Lock()
	if d.mu.index == len(d.mu.deck) {
		d.rng.Shuffle(len(d.mu.deck), func(i, j int) {
			d.mu.deck[i], d.mu.deck[j] = d.mu.deck[j], d.mu.deck[i]
		})
		d.mu.index = 0
	}
	result := d.mu.deck[d.mu.index]
	d.mu.index++
	d.mu.Unlock()
	return result
}