forked from ava-labs/avalanchego
-
Notifications
You must be signed in to change notification settings - Fork 4
/
rand.go
97 lines (83 loc) · 2.36 KB
/
rand.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package sampler
import (
"math"
"sync"
"time"
"gonum.org/v1/gonum/mathext/prng"
)
var globalRNG = newRNG()
func newRNG() *rng {
// We don't use a cryptographically secure source of randomness here, as
// there's no need to ensure a truly random sampling.
source := prng.NewMT19937()
source.Seed(uint64(time.Now().UnixNano()))
return &rng{rng: source}
}
func Seed(seed int64) {
globalRNG.Seed(seed)
}
type source interface {
Seed(uint64)
Uint64() uint64
}
type rng struct {
lock sync.Mutex
rng source
}
// Seed uses the provided seed value to initialize the generator to a
// deterministic state.
func (r *rng) Seed(seed int64) {
r.lock.Lock()
r.rng.Seed(uint64(seed))
r.lock.Unlock()
}
// Uint64Inclusive returns a pseudo-random number in [0,n].
//
// Invariant: The result of this function is stored in chain state, so any
// modifications are considered breaking.
func (r *rng) Uint64Inclusive(n uint64) uint64 {
switch {
// n+1 is power of two, so we can just mask
//
// Note: This does work for MaxUint64 as overflow is explicitly part of the
// compiler specification: https://go.dev/ref/spec#Integer_overflow
case n&(n+1) == 0:
return r.uint64() & n
// n is greater than MaxUint64/2 so we need to just iterate until we get a
// number in the requested range.
case n > math.MaxInt64:
v := r.uint64()
for v > n {
v = r.uint64()
}
return v
// n is less than MaxUint64/2 so we generate a number in the range
// [0, k*(n+1)) where k is the largest integer such that k*(n+1) is less
// than or equal to MaxUint64/2. We can't easily find k such that k*(n+1) is
// less than or equal to MaxUint64 because the calculation would overflow.
//
// ref: https://github.com/golang/go/blob/ce10e9d84574112b224eae88dc4e0f43710808de/src/math/rand/rand.go#L127-L132
default:
max := (1 << 63) - 1 - (1<<63)%(n+1)
v := r.uint63()
for v > max {
v = r.uint63()
}
return v % (n + 1)
}
}
// uint63 returns a random number in [0, MaxInt64]
func (r *rng) uint63() uint64 {
return r.uint64() & math.MaxInt64
}
// uint64 returns a random number in [0, MaxUint64]
func (r *rng) uint64() uint64 {
// Note: We must grab a write lock here because rng.Uint64 internally
// modifies state.
r.lock.Lock()
n := r.rng.Uint64()
r.lock.Unlock()
return n
}