From faad082efee3752fab2147c6b18124b8d13d3838 Mon Sep 17 00:00:00 2001 From: Vasiliy Tolstov Date: Wed, 15 Sep 2021 17:51:25 +0300 Subject: [PATCH] util/rand: add Shuffle func Signed-off-by: Vasiliy Tolstov --- util/rand/rand.go | 34 +++++++++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/util/rand/rand.go b/util/rand/rand.go index 5496d521..4b0b5ceb 100644 --- a/util/rand/rand.go +++ b/util/rand/rand.go @@ -1,17 +1,17 @@ package rand import ( - "crypto/rand" + crand "crypto/rand" "encoding/binary" ) -// Rand is a wrapper around crypto/rand that adds some convenience functions known from math/rand. +// Rand is a wrapper around crypto/rand that adds some convenience functions known from math/rand type Rand struct { buf [8]byte } func (r *Rand) Int31() int32 { - _, _ = rand.Read(r.buf[:4]) + _, _ = crand.Read(r.buf[:4]) return int32(binary.BigEndian.Uint32(r.buf[:4]) & ^uint32(1<<31)) } @@ -54,11 +54,11 @@ func (r *Rand) Intn(n int) int { } func (r *Rand) Int63() int64 { - _, _ = rand.Read(r.buf[:]) + _, _ = crand.Read(r.buf[:]) return int64(binary.BigEndian.Uint64(r.buf[:]) & ^uint64(1<<63)) } -// copied from the standard library math/rand implementation of Int63n +// Int31n copied from the standard library math/rand implementation of Int31n func (r *Rand) Int31n(n int32) int32 { if n&(n-1) == 0 { // n is power of two, can mask return r.Int31() & (n - 1) @@ -71,6 +71,7 @@ func (r *Rand) Int31n(n int32) int32 { return v % n } +// Int63n copied from the standard library math/rand implementation of Int63n func (r *Rand) Int63n(n int64) int64 { if n&(n-1) == 0 { // n is power of two, can mask return r.Int63() & (n - 1) @@ -82,3 +83,26 @@ func (r *Rand) Int63n(n int64) int64 { } return v % n } + +// Shuffle copied from the standard library math/rand implementation of Shuffle +func (r *Rand) Shuffle(n int, swap func(i, j int)) { + if n < 0 { + panic("invalid argument to Shuffle") + } + + // Fisher-Yates shuffle: https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle + // Shuffle really ought not be called with n that doesn't fit in 32 bits. + // Not only will it take a very long time, but with 2³¹! possible permutations, + // there's no way that any PRNG can have a big enough internal state to + // generate even a minuscule percentage of the possible permutations. + // Nevertheless, the right API signature accepts an int n, so handle it as best we can. + i := n - 1 + for ; i > 1<<31-1-1; i-- { + j := int(r.Int63n(int64(i + 1))) + swap(i, j) + } + for ; i > 0; i-- { + j := int(r.Int31n(int32(i + 1))) + swap(i, j) + } +}