atomic_splitmix64.go

package fastrand

import (
	"math/rand"
	"sync/atomic"
)

type SplitMix64 uint64
var _ rand.Source64 = &SplitMix64{}

func (a *SplitMix64) Uint64() uint64 {
	// SplitMix64
	// http://xoshiro.di.unimi.it/splitmix64.c
	x := atomic.AddUint64((*uint64)(a), 0x9e3779b97f4a7c15)
	x ^= x >> 30
	x *= 0xbf58476d1ce4e5b9
	x ^= x >> 27
	x *= 0x94d049bb133111eb
	x ^= x >> 31
	return x
}

func (a *SplitMix64) Int63() int64 {
	return int64(a.Uint64() >> 1)
}

func (a *SplitMix64) Seed(s int64) {
	*a = uint64(s)
}

func (a *SplitMix64) Uint64n(n uint64) uint64 {
	// random_bounded
	// https://lemire.me/blog/2016/06/30/fast-random-shuffling/
	x := a.Uint64()
	h, l := bits.Mul64(x, n)
	if l < n {
		t := (^n) % n
		for l < t {
			x = a.Uint64()
			h, l = bits.Mul64(x, n)
		}
	}
	return h
}

type SplitMix64Bool struct {
	b SplitMix64
	t uint64
}

func NewProb(prob float64) {
	if prob < 0 || prob > 1 {
		panic("invalid probability")
	}

	x := new(big.Rat).Mul(
		new(big.Rat).SetFloat64(prob), 
		new(big.Rat).SetUint64(math.MaxUint64),
	)
	t := new(big.Float).SetRat(x).Uint64()
	return SplitMix64Bool{t: t}
}

func NextBool() bool {
	return t.b.Uint64() <= t.t
}