rand: allow serialization of PRNG state
Fixes golang/go#35886.
Change-Id: If46dfc2bcbf4a7db84cb66c57f2d377c49d2efc0
Reviewed-on: https://go-review.googlesource.com/c/exp/+/209377
Reviewed-by: Rob Pike <r@golang.org>
diff --git a/rand/rand_test.go b/rand/rand_test.go
index 545b7dd..4383e4f 100644
--- a/rand/rand_test.go
+++ b/rand/rand_test.go
@@ -14,6 +14,7 @@
"runtime"
"testing"
"testing/iotest"
+ "time"
)
const (
@@ -457,6 +458,34 @@
}
}
+func TestPCGSourceRoundTrip(t *testing.T) {
+ var src PCGSource
+ src.Seed(uint64(time.Now().Unix()))
+
+ src.Uint64() // Step PRNG once to makes sure high and low are different.
+
+ buf, err := src.MarshalBinary()
+ if err != nil {
+ t.Errorf("unexpected error marshaling state: %v", err)
+ }
+
+ var dst PCGSource
+ // Get dst into a non-zero state.
+ dst.Seed(1)
+ for i := 0; i < 10; i++ {
+ dst.Uint64()
+ }
+
+ err = dst.UnmarshalBinary(buf)
+ if err != nil {
+ t.Errorf("unexpected error unmarshaling state: %v", err)
+ }
+
+ if dst != src {
+ t.Errorf("mismatch between generator states: got:%+v want:%+v", dst, src)
+ }
+}
+
// Benchmarks
func BenchmarkSource(b *testing.B) {
diff --git a/rand/rng.go b/rand/rng.go
index f4b0e19..17cee10 100644
--- a/rand/rng.go
+++ b/rand/rng.go
@@ -4,7 +4,11 @@
package rand
-import "math/bits"
+import (
+ "encoding/binary"
+ "io"
+ "math/bits"
+)
// PCGSource is an implementation of a 64-bit permuted congruential
// generator as defined in
@@ -67,3 +71,21 @@
pcg.low = lo
pcg.high = hi
}
+
+// MarshalBinary returns the binary representation of the current state of the generator.
+func (pcg *PCGSource) MarshalBinary() ([]byte, error) {
+ var buf [16]byte
+ binary.BigEndian.PutUint64(buf[:8], pcg.high)
+ binary.BigEndian.PutUint64(buf[8:], pcg.low)
+ return buf[:], nil
+}
+
+// UnmarshalBinary sets the state of the generator to the state represented in data.
+func (pcg *PCGSource) UnmarshalBinary(data []byte) error {
+ if len(data) < 16 {
+ return io.ErrUnexpectedEOF
+ }
+ pcg.low = binary.BigEndian.Uint64(data[8:])
+ pcg.high = binary.BigEndian.Uint64(data[:8])
+ return nil
+}
