[gostls13.git] / src / math / rand / v2 / rand.go

// Copyright 2009 The Go 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 rand implements pseudo-random number generators suitable for tasks
// such as simulation, but it should not be used for security-sensitive work.
//
// Random numbers are generated by a [Source], usually wrapped in a [Rand].
// Both types should be used by a single goroutine at a time: sharing among
// multiple goroutines requires some kind of synchronization.
//
// Top-level functions, such as [Float64] and [Int],
// are safe for concurrent use by multiple goroutines.
//
// This package's outputs might be easily predictable regardless of how it's
// seeded. For random numbers suitable for security-sensitive work, see the
// crypto/rand package.
package rand

import (
        "math/bits"
        _ "unsafe" // for go:linkname
)

// A Source is a source of uniformly-distributed
// pseudo-random uint64 values in the range [0, 1<<64).
//
// A Source is not safe for concurrent use by multiple goroutines.
type Source interface {
        Uint64() uint64
}

// NewSource returns a new pseudo-random Source seeded with the given value.
// Unlike the default Source used by top-level functions, this source is not
// safe for concurrent use by multiple goroutines.
// The returned Source implements Source64.
func NewSource(seed int64) Source {
        return newSource(seed)
}

func newSource(seed int64) *rngSource {
        var rng rngSource
        rng.Seed(seed)
        return &rng
}

// A Rand is a source of random numbers.
type Rand struct {
        src Source
}

// New returns a new Rand that uses random values from src
// to generate other random values.
func New(src Source) *Rand {
        return &Rand{src: src}
}

// Int64 returns a non-negative pseudo-random 63-bit integer as an int64.
func (r *Rand) Int64() int64 { return int64(r.src.Uint64() &^ (1 << 63)) }

// Uint32 returns a pseudo-random 32-bit value as a uint32.
func (r *Rand) Uint32() uint32 { return uint32(r.src.Uint64() >> 32) }

// Uint64 returns a pseudo-random 64-bit value as a uint64.
func (r *Rand) Uint64() uint64 { return r.src.Uint64() }

// Int32 returns a non-negative pseudo-random 31-bit integer as an int32.
func (r *Rand) Int32() int32 { return int32(r.src.Uint64() >> 33) }

// Int returns a non-negative pseudo-random int.
func (r *Rand) Int() int { return int(uint(r.src.Uint64()) << 1 >> 1) }

// Int64N returns, as an int64, a non-negative pseudo-random number in the half-open interval [0,n).
// It panics if n <= 0.
func (r *Rand) Int64N(n int64) int64 {
        if n <= 0 {
                panic("invalid argument to Int64N")
        }
        return int64(r.uint64n(uint64(n)))
}

// Uint64N returns, as a uint64, a non-negative pseudo-random number in the half-open interval [0,n).
// It panics if n == 0.
func (r *Rand) Uint64N(n uint64) uint64 {
        if n == 0 {
                panic("invalid argument to Uint64N")
        }
        return r.uint64n(n)
}

// uint64n is the no-bounds-checks version of Uint64N.
func (r *Rand) uint64n(n uint64) uint64 {
        if is32bit && uint64(uint32(n)) == n {
                return uint64(r.uint32n(uint32(n)))
        }
        if n&(n-1) == 0 { // n is power of two, can mask
                return r.Uint64() & (n - 1)
        }

        // Suppose we have a uint64 x uniform in the range [0,2⁶⁴)
        // and want to reduce it to the range [0,n) preserving exact uniformity.
        // We can simulate a scaling arbitrary precision x * (n/2⁶⁴) by
        // the high bits of a double-width multiply of x*n, meaning (x*n)/2⁶⁴.
        // Since there are 2⁶⁴ possible inputs x and only n possible outputs,
        // the output is necessarily biased if n does not divide 2⁶⁴.
        // In general (x*n)/2⁶⁴ = k for x*n in [k*2⁶⁴,(k+1)*2⁶⁴).
        // There are either floor(2⁶⁴/n) or ceil(2⁶⁴/n) possible products
        // in that range, depending on k.
        // But suppose we reject the sample and try again when
        // x*n is in [k*2⁶⁴, k*2⁶⁴+(2⁶⁴%n)), meaning rejecting fewer than n possible
        // outcomes out of the 2⁶⁴.
        // Now there are exactly floor(2⁶⁴/n) possible ways to produce
        // each output value k, so we've restored uniformity.
        // To get valid uint64 math, 2⁶⁴ % n = (2⁶⁴ - n) % n = -n % n,
        // so the direct implementation of this algorithm would be:
        //
        //      hi, lo := bits.Mul64(r.Uint64(), n)
        //      thresh := -n % n
        //      for lo < thresh {
        //              hi, lo = bits.Mul64(r.Uint64(), n)
        //      }
        //
        // That still leaves an expensive 64-bit division that we would rather avoid.
        // We know that thresh < n, and n is usually much less than 2⁶⁴, so we can
        // avoid the last four lines unless lo < n.
        //
        // See also:
        // https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction
        // https://lemire.me/blog/2016/06/30/fast-random-shuffling
        hi, lo := bits.Mul64(r.Uint64(), n)
        if lo < n {
                thresh := -n % n
                for lo < thresh {
                        hi, lo = bits.Mul64(r.Uint64(), n)
                }
        }
        return hi
}

// uint32n is an identical computation to uint64n
// but optimized for 32-bit systems.
func (r *Rand) uint32n(n uint32) uint32 {
        if n&(n-1) == 0 { // n is power of two, can mask
                return uint32(r.Uint64()) & (n - 1)
        }
        // On 64-bit systems we still use the uint64 code below because
        // the probability of a random uint64 lo being < a uint32 n is near zero,
        // meaning the unbiasing loop almost never runs.
        // On 32-bit systems, here we need to implement that same logic in 32-bit math,
        // both to preserve the exact output sequence observed on 64-bit machines
        // and to preserve the optimization that the unbiasing loop almost never runs.
        //
        // We want to compute
        //      hi, lo := bits.Mul64(r.Uint64(), n)
        // In terms of 32-bit halves, this is:
        //      x1:x0 := r.Uint64()
        //      0:hi, lo1:lo0 := bits.Mul64(x1:x0, 0:n)
        // Writing out the multiplication in terms of bits.Mul32 allows
        // using direct hardware instructions and avoiding
        // the computations involving these zeros.
        x := r.Uint64()
        lo1a, lo0 := bits.Mul32(uint32(x), n)
        hi, lo1b := bits.Mul32(uint32(x>>32), n)
        lo1, c := bits.Add32(lo1a, lo1b, 0)
        hi += c
        if lo1 == 0 && lo0 < uint32(n) {
                n64 := uint64(n)
                thresh := uint32(-n64 % n64)
                for lo1 == 0 && lo0 < thresh {
                        x := r.Uint64()
                        lo1a, lo0 = bits.Mul32(uint32(x), n)
                        hi, lo1b = bits.Mul32(uint32(x>>32), n)
                        lo1, c = bits.Add32(lo1a, lo1b, 0)
                        hi += c
                }
        }
        return hi
}

// Int32N returns, as an int32, a non-negative pseudo-random number in the half-open interval [0,n).
// It panics if n <= 0.
func (r *Rand) Int32N(n int32) int32 {
        if n <= 0 {
                panic("invalid argument to Int32N")
        }
        return int32(r.uint64n(uint64(n)))
}

// Uint32N returns, as a uint32, a non-negative pseudo-random number in the half-open interval [0,n).
// It panics if n == 0.
func (r *Rand) Uint32N(n uint32) uint32 {
        if n == 0 {
                panic("invalid argument to Uint32N")
        }
        return uint32(r.uint64n(uint64(n)))
}

const is32bit = ^uint(0)>>32 == 0

// IntN returns, as an int, a non-negative pseudo-random number in the half-open interval [0,n).
// It panics if n <= 0.
func (r *Rand) IntN(n int) int {
        if n <= 0 {
                panic("invalid argument to IntN")
        }
        return int(r.uint64n(uint64(n)))
}

// UintN returns, as a uint, a non-negative pseudo-random number in the half-open interval [0,n).
// It panics if n == 0.
func (r *Rand) UintN(n uint) uint {
        if n == 0 {
                panic("invalid argument to UintN")
        }
        return uint(r.uint64n(uint64(n)))
}

// Float64 returns, as a float64, a pseudo-random number in the half-open interval [0.0,1.0).
func (r *Rand) Float64() float64 {
        // There are exactly 1<<53 float64s in [0,1). Use Intn(1<<53) / (1<<53).
        return float64(r.Uint64()<<11>>11) / (1 << 53)
}

// Float32 returns, as a float32, a pseudo-random number in the half-open interval [0.0,1.0).
func (r *Rand) Float32() float32 {
        // There are exactly 1<<24 float32s in [0,1). Use Intn(1<<24) / (1<<24).
        return float32(r.Uint32()<<8>>8) / (1 << 24)
}

// Perm returns, as a slice of n ints, a pseudo-random permutation of the integers
// in the half-open interval [0,n).
func (r *Rand) Perm(n int) []int {
        p := make([]int, n)
        for i := range p {
                p[i] = i
        }
        r.Shuffle(len(p), func(i, j int) { p[i], p[j] = p[j], p[i] })
        return p
}

// Shuffle pseudo-randomizes the order of elements.
// n is the number of elements. Shuffle panics if n < 0.
// swap swaps the elements with indexes i and j.
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.
        for i := n - 1; i > 0; i-- {
                j := int(r.uint64n(uint64(i + 1)))
                swap(i, j)
        }
}

/*
 * Top-level convenience functions
 */

// globalRand is the source of random numbers for the top-level
// convenience functions.
var globalRand = &Rand{src: &fastSource{}}

//go:linkname fastrand64
func fastrand64() uint64

// fastSource is a Source that uses the runtime fastrand functions.
type fastSource struct{}

func (*fastSource) Int64() int64 {
        return int64(fastrand64() & rngMask)
}

func (*fastSource) Uint64() uint64 {
        return fastrand64()
}

// Int64 returns a non-negative pseudo-random 63-bit integer as an int64
// from the default Source.
func Int64() int64 { return globalRand.Int64() }

// Uint32 returns a pseudo-random 32-bit value as a uint32
// from the default Source.
func Uint32() uint32 { return globalRand.Uint32() }

// Uint64N returns, as a uint64, a pseudo-random number in the half-open interval [0,n)
// from the default Source.
// It panics if n <= 0.
func Uint64N(n uint64) uint64 { return globalRand.Uint64N(n) }

// Uint32N returns, as a uint32, a pseudo-random number in the half-open interval [0,n)
// from the default Source.
// It panics if n <= 0.
func Uint32N(n uint32) uint32 { return globalRand.Uint32N(n) }

// Uint64 returns a pseudo-random 64-bit value as a uint64
// from the default Source.
func Uint64() uint64 { return globalRand.Uint64() }

// Int32 returns a non-negative pseudo-random 31-bit integer as an int32
// from the default Source.
func Int32() int32 { return globalRand.Int32() }

// Int returns a non-negative pseudo-random int from the default Source.
func Int() int { return globalRand.Int() }

// Int64N returns, as an int64, a pseudo-random number in the half-open interval [0,n)
// from the default Source.
// It panics if n <= 0.
func Int64N(n int64) int64 { return globalRand.Int64N(n) }

// Int32N returns, as an int32, a pseudo-random number in the half-open interval [0,n)
// from the default Source.
// It panics if n <= 0.
func Int32N(n int32) int32 { return globalRand.Int32N(n) }

// IntN returns, as an int, a pseudo-random number in the half-open interval [0,n)
// from the default Source.
// It panics if n <= 0.
func IntN(n int) int { return globalRand.IntN(n) }

// UintN returns, as a uint, a pseudo-random number in the half-open interval [0,n)
// from the default Source.
// It panics if n <= 0.
func UintN(n uint) uint { return globalRand.UintN(n) }

// N returns a pseudo-random number in the half-open interval [0,n) from the default Source.
// The type parameter Int can be any integer type.
// It panics if n <= 0.
func N[Int intType](n Int) Int {
        if n <= 0 {
                panic("invalid argument to N")
        }
        return Int(globalRand.uint64n(uint64(n)))
}

type intType interface {
        ~int | ~int8 | ~int16 | ~int32 | ~int64 |
                ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr
}

// Float64 returns, as a float64, a pseudo-random number in the half-open interval [0.0,1.0)
// from the default Source.
func Float64() float64 { return globalRand.Float64() }

// Float32 returns, as a float32, a pseudo-random number in the half-open interval [0.0,1.0)
// from the default Source.
func Float32() float32 { return globalRand.Float32() }

// Perm returns, as a slice of n ints, a pseudo-random permutation of the integers
// in the half-open interval [0,n) from the default Source.
func Perm(n int) []int { return globalRand.Perm(n) }

// Shuffle pseudo-randomizes the order of elements using the default Source.
// n is the number of elements. Shuffle panics if n < 0.
// swap swaps the elements with indexes i and j.
func Shuffle(n int, swap func(i, j int)) { globalRand.Shuffle(n, swap) }

// NormFloat64 returns a normally distributed float64 in the range
// [-math.MaxFloat64, +math.MaxFloat64] with
// standard normal distribution (mean = 0, stddev = 1)
// from the default Source.
// To produce a different normal distribution, callers can
// adjust the output using:
//
//      sample = NormFloat64() * desiredStdDev + desiredMean
func NormFloat64() float64 { return globalRand.NormFloat64() }

// ExpFloat64 returns an exponentially distributed float64 in the range
// (0, +math.MaxFloat64] with an exponential distribution whose rate parameter
// (lambda) is 1 and whose mean is 1/lambda (1) from the default Source.
// To produce a distribution with a different rate parameter,
// callers can adjust the output using:
//
//      sample = ExpFloat64() / desiredRateParameter
func ExpFloat64() float64 { return globalRand.ExpFloat64() }