pcg.go

     1  // Copyright 2023 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package rand
     6  
     7  import (
     8  	"errors"
     9  	"math/bits"
    10  )
    11  
    12  // https://numpy.org/devdocs/reference/random/upgrading-pcg64.html
    13  // https://github.com/imneme/pcg-cpp/commit/871d0494ee9c9a7b7c43f753e3d8ca47c26f8005
    14  
    15  // A PCG is a PCG generator with 128 bits of internal state.
    16  // A zero PCG is equivalent to NewPCG(0, 0).
    17  type PCG struct {
    18  	hi uint64
    19  	lo uint64
    20  }
    21  
    22  // NewPCG returns a new PCG seeded with the given values.
    23  func NewPCG(seed1, seed2 uint64) *PCG {
    24  	return &PCG{seed1, seed2}
    25  }
    26  
    27  // Seed resets the PCG to behave the same way as NewPCG(seed1, seed2).
    28  func (p *PCG) Seed(seed1, seed2 uint64) {
    29  	p.hi = seed1
    30  	p.lo = seed2
    31  }
    32  
    33  // binary.bigEndian.Uint64, copied to avoid dependency
    34  func beUint64(b []byte) uint64 {
    35  	_ = b[7] // bounds check hint to compiler; see golang.org/issue/14808
    36  	return uint64(b[7]) | uint64(b[6])<<8 | uint64(b[5])<<16 | uint64(b[4])<<24 |
    37  		uint64(b[3])<<32 | uint64(b[2])<<40 | uint64(b[1])<<48 | uint64(b[0])<<56
    38  }
    39  
    40  // binary.bigEndian.PutUint64, copied to avoid dependency
    41  func bePutUint64(b []byte, v uint64) {
    42  	_ = b[7] // early bounds check to guarantee safety of writes below
    43  	b[0] = byte(v >> 56)
    44  	b[1] = byte(v >> 48)
    45  	b[2] = byte(v >> 40)
    46  	b[3] = byte(v >> 32)
    47  	b[4] = byte(v >> 24)
    48  	b[5] = byte(v >> 16)
    49  	b[6] = byte(v >> 8)
    50  	b[7] = byte(v)
    51  }
    52  
    53  // MarshalBinary implements the encoding.BinaryMarshaler interface.
    54  func (p *PCG) MarshalBinary() ([]byte, error) {
    55  	b := make([]byte, 20)
    56  	copy(b, "pcg:")
    57  	bePutUint64(b[4:], p.hi)
    58  	bePutUint64(b[4+8:], p.lo)
    59  	return b, nil
    60  }
    61  
    62  var errUnmarshalPCG = errors.New("invalid PCG encoding")
    63  
    64  // UnmarshalBinary implements the encoding.BinaryUnmarshaler interface.
    65  func (p *PCG) UnmarshalBinary(data []byte) error {
    66  	if len(data) != 20 || string(data[:4]) != "pcg:" {
    67  		return errUnmarshalPCG
    68  	}
    69  	p.hi = beUint64(data[4:])
    70  	p.lo = beUint64(data[4+8:])
    71  	return nil
    72  }
    73  
    74  func (p *PCG) next() (hi, lo uint64) {
    75  	// https://github.com/imneme/pcg-cpp/blob/428802d1a5/include/pcg_random.hpp#L161
    76  	//
    77  	// Numpy's PCG multiplies by the 64-bit value cheapMul
    78  	// instead of the 128-bit value used here and in the official PCG code.
    79  	// This does not seem worthwhile, at least for Go: not having any high
    80  	// bits in the multiplier reduces the effect of low bits on the highest bits,
    81  	// and it only saves 1 multiply out of 3.
    82  	// (On 32-bit systems, it saves 1 out of 6, since Mul64 is doing 4.)
    83  	const (
    84  		mulHi = 2549297995355413924
    85  		mulLo = 4865540595714422341
    86  		incHi = 6364136223846793005
    87  		incLo = 1442695040888963407
    88  	)
    89  
    90  	// state = state * mul + inc
    91  	hi, lo = bits.Mul64(p.lo, mulLo)
    92  	hi += p.hi*mulLo + p.lo*mulHi
    93  	lo, c := bits.Add64(lo, incLo, 0)
    94  	hi, _ = bits.Add64(hi, incHi, c)
    95  	p.lo = lo
    96  	p.hi = hi
    97  	return hi, lo
    98  }
    99  
   100  // Uint64 return a uniformly-distributed random uint64 value.
   101  func (p *PCG) Uint64() uint64 {
   102  	hi, lo := p.next()
   103  
   104  	// XSL-RR would be
   105  	//	hi, lo := p.next()
   106  	//	return bits.RotateLeft64(lo^hi, -int(hi>>58))
   107  	// but Numpy uses DXSM and O'Neill suggests doing the same.
   108  	// See https://github.com/golang/go/issues/21835#issuecomment-739065688
   109  	// and following comments.
   110  
   111  	// DXSM "double xorshift multiply"
   112  	// https://github.com/imneme/pcg-cpp/blob/428802d1a5/include/pcg_random.hpp#L1015
   113  
   114  	// https://github.com/imneme/pcg-cpp/blob/428802d1a5/include/pcg_random.hpp#L176
   115  	const cheapMul = 0xda942042e4dd58b5
   116  	hi ^= hi >> 32
   117  	hi *= cheapMul
   118  	hi ^= hi >> 48
   119  	hi *= (lo | 1)
   120  	return hi
   121  }
   122
View as plain text