util/rand: add Shuffle func

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
This commit is contained in:
Василий Толстов 2021-09-15 17:51:25 +03:00
parent 8ab35cbd9b
commit faad082efe

View File

@ -1,17 +1,17 @@
package rand package rand
import ( import (
"crypto/rand" crand "crypto/rand"
"encoding/binary" "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 { type Rand struct {
buf [8]byte buf [8]byte
} }
func (r *Rand) Int31() int32 { 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)) 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 { func (r *Rand) Int63() int64 {
_, _ = rand.Read(r.buf[:]) _, _ = crand.Read(r.buf[:])
return int64(binary.BigEndian.Uint64(r.buf[:]) & ^uint64(1<<63)) 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 { func (r *Rand) Int31n(n int32) int32 {
if n&(n-1) == 0 { // n is power of two, can mask if n&(n-1) == 0 { // n is power of two, can mask
return r.Int31() & (n - 1) return r.Int31() & (n - 1)
@ -71,6 +71,7 @@ func (r *Rand) Int31n(n int32) int32 {
return v % n return v % n
} }
// Int63n copied from the standard library math/rand implementation of Int63n
func (r *Rand) Int63n(n int64) int64 { func (r *Rand) Int63n(n int64) int64 {
if n&(n-1) == 0 { // n is power of two, can mask if n&(n-1) == 0 { // n is power of two, can mask
return r.Int63() & (n - 1) return r.Int63() & (n - 1)
@ -82,3 +83,26 @@ func (r *Rand) Int63n(n int64) int64 {
} }
return v % n 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)
}
}