1818package rand
1919
2020import (
21+ "math/bits"
2122 _ "unsafe" // for go:linkname
2223)
2324
@@ -58,37 +59,122 @@ func New(src Source) *Rand {
5859func (r * Rand ) Int64 () int64 { return int64 (r .src .Uint64 () &^ (1 << 63 )) }
5960
6061// Uint32 returns a pseudo-random 32-bit value as a uint32.
61- func (r * Rand ) Uint32 () uint32 { return uint32 (r .Int64 () >> 31 ) }
62+ func (r * Rand ) Uint32 () uint32 { return uint32 (r .src . Uint64 () >> 32 ) }
6263
6364// Uint64 returns a pseudo-random 64-bit value as a uint64.
64- func (r * Rand ) Uint64 () uint64 {
65- return r .src .Uint64 ()
66- }
65+ func (r * Rand ) Uint64 () uint64 { return r .src .Uint64 () }
6766
6867// Int32 returns a non-negative pseudo-random 31-bit integer as an int32.
69- func (r * Rand ) Int32 () int32 { return int32 (r .Int64 () >> 32 ) }
68+ func (r * Rand ) Int32 () int32 { return int32 (r .src . Uint64 () >> 33 ) }
7069
7170// Int returns a non-negative pseudo-random int.
72- func (r * Rand ) Int () int {
73- u := uint (r .Int64 ())
74- return int (u << 1 >> 1 ) // clear sign bit if int == int32
75- }
71+ func (r * Rand ) Int () int { return int (uint (r .src .Uint64 ()) << 1 >> 1 ) }
7672
7773// Int64N returns, as an int64, a non-negative pseudo-random number in the half-open interval [0,n).
7874// It panics if n <= 0.
7975func (r * Rand ) Int64N (n int64 ) int64 {
8076 if n <= 0 {
8177 panic ("invalid argument to Int64N" )
8278 }
79+ return int64 (r .uint64n (uint64 (n )))
80+ }
81+
82+ // Uint64N returns, as a uint64, a non-negative pseudo-random number in the half-open interval [0,n).
83+ // It panics if n == 0.
84+ func (r * Rand ) Uint64N (n uint64 ) uint64 {
85+ if n == 0 {
86+ panic ("invalid argument to Uint64N" )
87+ }
88+ return r .uint64n (n )
89+ }
90+
91+ // uint64n is the no-bounds-checks version of Uint64N.
92+ func (r * Rand ) uint64n (n uint64 ) uint64 {
93+ if is32bit && uint64 (uint32 (n )) == n {
94+ return uint64 (r .uint32n (uint32 (n )))
95+ }
8396 if n & (n - 1 ) == 0 { // n is power of two, can mask
84- return r .Int64 () & (n - 1 )
97+ return r .Uint64 () & (n - 1 )
8598 }
86- max := int64 ((1 << 63 ) - 1 - (1 << 63 )% uint64 (n ))
87- v := r .Int64 ()
88- for v > max {
89- v = r .Int64 ()
99+
100+ // Suppose we have a uint64 x uniform in the range [0,2⁶⁴)
101+ // and want to reduce it to the range [0,n) preserving exact uniformity.
102+ // We can simulate a scaling arbitrary precision x * (n/2⁶⁴) by
103+ // the high bits of a double-width multiply of x*n, meaning (x*n)/2⁶⁴.
104+ // Since there are 2⁶⁴ possible inputs x and only n possible outputs,
105+ // the output is necessarily biased if n does not divide 2⁶⁴.
106+ // In general (x*n)/2⁶⁴ = k for x*n in [k*2⁶⁴,(k+1)*2⁶⁴).
107+ // There are either floor(2⁶⁴/n) or ceil(2⁶⁴/n) possible products
108+ // in that range, depending on k.
109+ // But suppose we reject the sample and try again when
110+ // x*n is in [k*2⁶⁴, k*2⁶⁴+(2⁶⁴%n)), meaning rejecting fewer than n possible
111+ // outcomes out of the 2⁶⁴.
112+ // Now there are exactly floor(2⁶⁴/n) possible ways to produce
113+ // each output value k, so we've restored uniformity.
114+ // To get valid uint64 math, 2⁶⁴ % n = (2⁶⁴ - n) % n = -n % n,
115+ // so the direct implementation of this algorithm would be:
116+ //
117+ // hi, lo := bits.Mul64(r.Uint64(), n)
118+ // thresh := -n % n
119+ // for lo < thresh {
120+ // hi, lo = bits.Mul64(r.Uint64(), n)
121+ // }
122+ //
123+ // That still leaves an expensive 64-bit division that we would rather avoid.
124+ // We know that thresh < n, and n is usually much less than 2⁶⁴, so we can
125+ // avoid the last four lines unless lo < n.
126+ //
127+ // See also:
128+ // https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction
129+ // https://lemire.me/blog/2016/06/30/fast-random-shuffling
130+ hi , lo := bits .Mul64 (r .Uint64 (), n )
131+ if lo < n {
132+ thresh := - n % n
133+ for lo < thresh {
134+ hi , lo = bits .Mul64 (r .Uint64 (), n )
135+ }
90136 }
91- return v % n
137+ return hi
138+ }
139+
140+ // uint32n is an identical computation to uint64n
141+ // but optimized for 32-bit systems.
142+ func (r * Rand ) uint32n (n uint32 ) uint32 {
143+ if n & (n - 1 ) == 0 { // n is power of two, can mask
144+ return uint32 (r .Uint64 ()) & (n - 1 )
145+ }
146+ // On 64-bit systems we still use the uint64 code below because
147+ // the probability of a random uint64 lo being < a uint32 n is near zero,
148+ // meaning the unbiasing loop almost never runs.
149+ // On 32-bit systems, here we need to implement that same logic in 32-bit math,
150+ // both to preserve the exact output sequence observed on 64-bit machines
151+ // and to preserve the optimization that the unbiasing loop almost never runs.
152+ //
153+ // We want to compute
154+ // hi, lo := bits.Mul64(r.Uint64(), n)
155+ // In terms of 32-bit halves, this is:
156+ // x1:x0 := r.Uint64()
157+ // 0:hi, lo1:lo0 := bits.Mul64(x1:x0, 0:n)
158+ // Writing out the multiplication in terms of bits.Mul32 allows
159+ // using direct hardware instructions and avoiding
160+ // the computations involving these zeros.
161+ x := r .Uint64 ()
162+ lo1a , lo0 := bits .Mul32 (uint32 (x ), n )
163+ hi , lo1b := bits .Mul32 (uint32 (x >> 32 ), n )
164+ lo1 , c := bits .Add32 (lo1a , lo1b , 0 )
165+ hi += c
166+ if lo1 == 0 && lo0 < uint32 (n ) {
167+ n64 := uint64 (n )
168+ thresh := uint32 (- n64 % n64 )
169+ for lo1 == 0 && lo0 < thresh {
170+ x := r .Uint64 ()
171+ lo1a , lo0 = bits .Mul32 (uint32 (x ), n )
172+ hi , lo1b = bits .Mul32 (uint32 (x >> 32 ), n )
173+ lo1 , c = bits .Add32 (lo1a , lo1b , 0 )
174+ hi += c
175+ }
176+ }
177+ return hi
92178}
93179
94180// Int32N returns, as an int32, a non-negative pseudo-random number in the half-open interval [0,n).
@@ -97,51 +183,36 @@ func (r *Rand) Int32N(n int32) int32 {
97183 if n <= 0 {
98184 panic ("invalid argument to Int32N" )
99185 }
100- if n & (n - 1 ) == 0 { // n is power of two, can mask
101- return r .Int32 () & (n - 1 )
102- }
103- max := int32 ((1 << 31 ) - 1 - (1 << 31 )% uint32 (n ))
104- v := r .Int32 ()
105- for v > max {
106- v = r .Int32 ()
107- }
108- return v % n
186+ return int32 (r .uint64n (uint64 (n )))
109187}
110188
111- // int31n returns, as an int32, a non-negative pseudo-random number in the half-open interval [0,n).
112- // n must be > 0, but int31n does not check this; the caller must ensure it.
113- // int31n exists because Int32N is inefficient, but Go 1 compatibility
114- // requires that the stream of values produced by math/rand/v2 remain unchanged.
115- // int31n can thus only be used internally, by newly introduced APIs.
116- //
117- // For implementation details, see:
118- // https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction
119- // https://lemire.me/blog/2016/06/30/fast-random-shuffling
120- func (r * Rand ) int31n (n int32 ) int32 {
121- v := r .Uint32 ()
122- prod := uint64 (v ) * uint64 (n )
123- low := uint32 (prod )
124- if low < uint32 (n ) {
125- thresh := uint32 (- n ) % uint32 (n )
126- for low < thresh {
127- v = r .Uint32 ()
128- prod = uint64 (v ) * uint64 (n )
129- low = uint32 (prod )
130- }
189+ // Uint32N returns, as a uint32, a non-negative pseudo-random number in the half-open interval [0,n).
190+ // It panics if n == 0.
191+ func (r * Rand ) Uint32N (n uint32 ) uint32 {
192+ if n == 0 {
193+ panic ("invalid argument to Uint32N" )
131194 }
132- return int32 ( prod >> 32 )
195+ return uint32 ( r . uint64n ( uint64 ( n )) )
133196}
134197
198+ const is32bit = ^ uint (0 )>> 32 == 0
199+
135200// IntN returns, as an int, a non-negative pseudo-random number in the half-open interval [0,n).
136201// It panics if n <= 0.
137202func (r * Rand ) IntN (n int ) int {
138203 if n <= 0 {
139204 panic ("invalid argument to IntN" )
140205 }
141- if n <= 1 << 31 - 1 {
142- return int (r .Int32N (int32 (n )))
206+ return int (r .uint64n (uint64 (n )))
207+ }
208+
209+ // UintN returns, as a uint, a non-negative pseudo-random number in the half-open interval [0,n).
210+ // It panics if n == 0.
211+ func (r * Rand ) UintN (n uint ) uint {
212+ if n == 0 {
213+ panic ("invalid argument to UintN" )
143214 }
144- return int (r .Int64N ( int64 (n )))
215+ return uint (r .uint64n ( uint64 (n )))
145216}
146217
147218// Float64 returns, as a float64, a pseudo-random number in the half-open interval [0.0,1.0).
@@ -214,13 +285,8 @@ func (r *Rand) Shuffle(n int, swap func(i, j int)) {
214285 // there's no way that any PRNG can have a big enough internal state to
215286 // generate even a minuscule percentage of the possible permutations.
216287 // Nevertheless, the right API signature accepts an int n, so handle it as best we can.
217- i := n - 1
218- for ; i > 1 << 31 - 1 - 1 ; i -- {
219- j := int (r .Int64N (int64 (i + 1 )))
220- swap (i , j )
221- }
222- for ; i > 0 ; i -- {
223- j := int (r .int31n (int32 (i + 1 )))
288+ for i := n - 1 ; i > 0 ; i -- {
289+ j := int (r .uint64n (uint64 (i + 1 )))
224290 swap (i , j )
225291 }
226292}
@@ -255,6 +321,16 @@ func Int64() int64 { return globalRand.Int64() }
255321// from the default Source.
256322func Uint32 () uint32 { return globalRand .Uint32 () }
257323
324+ // Uint64N returns, as a uint64, a pseudo-random number in the half-open interval [0,n)
325+ // from the default Source.
326+ // It panics if n <= 0.
327+ func Uint64N (n uint64 ) uint64 { return globalRand .Uint64N (n ) }
328+
329+ // Uint32N returns, as a uint32, a pseudo-random number in the half-open interval [0,n)
330+ // from the default Source.
331+ // It panics if n <= 0.
332+ func Uint32N (n uint32 ) uint32 { return globalRand .Uint32N (n ) }
333+
258334// Uint64 returns a pseudo-random 64-bit value as a uint64
259335// from the default Source.
260336func Uint64 () uint64 { return globalRand .Uint64 () }
@@ -266,21 +342,41 @@ func Int32() int32 { return globalRand.Int32() }
266342// Int returns a non-negative pseudo-random int from the default Source.
267343func Int () int { return globalRand .Int () }
268344
269- // Int64N returns, as an int64, a non-negative pseudo-random number in the half-open interval [0,n)
345+ // Int64N returns, as an int64, a pseudo-random number in the half-open interval [0,n)
270346// from the default Source.
271347// It panics if n <= 0.
272348func Int64N (n int64 ) int64 { return globalRand .Int64N (n ) }
273349
274- // Int32N returns, as an int32, a non-negative pseudo-random number in the half-open interval [0,n)
350+ // Int32N returns, as an int32, a pseudo-random number in the half-open interval [0,n)
275351// from the default Source.
276352// It panics if n <= 0.
277353func Int32N (n int32 ) int32 { return globalRand .Int32N (n ) }
278354
279- // IntN returns, as an int, a non-negative pseudo-random number in the half-open interval [0,n)
355+ // IntN returns, as an int, a pseudo-random number in the half-open interval [0,n)
280356// from the default Source.
281357// It panics if n <= 0.
282358func IntN (n int ) int { return globalRand .IntN (n ) }
283359
360+ // UintN returns, as a uint, a pseudo-random number in the half-open interval [0,n)
361+ // from the default Source.
362+ // It panics if n <= 0.
363+ func UintN (n uint ) uint { return globalRand .UintN (n ) }
364+
365+ // N returns a pseudo-random number in the half-open interval [0,n) from the default Source.
366+ // The type parameter Int can be any integer type.
367+ // It panics if n <= 0.
368+ func N [Int intType ](n Int ) Int {
369+ if n <= 0 {
370+ panic ("invalid argument to N" )
371+ }
372+ return Int (globalRand .uint64n (uint64 (n )))
373+ }
374+
375+ type intType interface {
376+ ~ int | ~ int8 | ~ int16 | ~ int32 | ~ int64 |
377+ ~ uint | ~ uint8 | ~ uint16 | ~ uint32 | ~ uint64 | ~ uintptr
378+ }
379+
284380// Float64 returns, as a float64, a pseudo-random number in the half-open interval [0.0,1.0)
285381// from the default Source.
286382func Float64 () float64 { return globalRand .Float64 () }
0 commit comments