v8.run/go/exp@v0.0.26-0.20230226010534-afcdbd3f782d/fastrand/rng.go (about) 1 package fastrand 2 3 import ( 4 "crypto/rand" 5 "encoding/binary" 6 "sync" 7 8 "v8.run/go/exp/fastrand/alg/splitmix64" 9 "v8.run/go/exp/util/noescape" 10 ) 11 12 type RNG struct { 13 state [2]uint64 14 } 15 16 func refill(r *RNG) { 17 var data [16]byte 18 _, err := noescape.Read(rand.Reader, data[:]) 19 if err != nil { 20 // If crypto/rand fails, we'll just use the runtime.fastrand implementation. 21 binary.LittleEndian.PutUint32(data[0:4], runtime_fastrand()) 22 binary.LittleEndian.PutUint32(data[4:8], runtime_fastrand()) 23 binary.LittleEndian.PutUint32(data[8:12], runtime_fastrand()) 24 binary.LittleEndian.PutUint32(data[12:16], runtime_fastrand()) 25 } 26 27 r.state[0] = binary.LittleEndian.Uint64(data[0:8]) 28 r.state[1] = binary.LittleEndian.Uint64(data[8:16]) 29 30 // Use Splitmix64 to initialize the state. 31 r.state[0] += 1757750930446974760 32 r.state[1] += 7151402297004559274 33 r.state[0] = splitmix64.Splitmix64(&r.state[0]) 34 r.state[1] = splitmix64.Splitmix64(&r.state[1]) 35 } 36 37 func newRNG() *RNG { 38 r := new(RNG) 39 refill(r) 40 return r 41 } 42 43 var rngPool sync.Pool = sync.Pool{ 44 New: func() interface{} { 45 return newRNG() 46 }, 47 } 48 49 func AcquireRNG() *RNG { 50 return rngPool.Get().(*RNG) 51 } 52 53 func ReleaseRNG(r *RNG) { 54 rngPool.Put(r) 55 } 56 57 func WithSeed(seed uint64) *RNG { 58 r := AcquireRNG() 59 r.SetSeed(seed) 60 return r 61 } 62 63 func (rng *RNG) SetSeed(seed uint64) { 64 rng.state[0] = splitmix64.Splitmix64(&seed) 65 rng.state[1] = splitmix64.Splitmix64(&seed) 66 } 67 68 // Release Put the RNG back into the pool. 69 // After calling this, the RNG is invalid and should not be used. 70 func (rng *RNG) Release() { 71 ReleaseRNG(rng) 72 } 73 74 // Refill Initialize the RNG with a new seed. 75 func (rng *RNG) Refill() { 76 refill(rng) 77 } 78 79 type FastRandReader struct { 80 RNG *RNG 81 } 82 83 func (r *FastRandReader) Read(p []byte) (n int, err error) { 84 if r.RNG == nil { 85 r.RNG = AcquireRNG() 86 } 87 n = len(p) 88 89 for len(p) >= 8 { 90 binary.LittleEndian.PutUint64(p, r.RNG.Uint64()) 91 p = p[8:] 92 } 93 94 if len(p) > 0 { 95 v := r.RNG.Uint64() 96 for i := 0; i < len(p); i++ { 97 v >>= 8 98 p[i] = byte(v) 99 } 100 } 101 102 return 103 }