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  }