math/rand: fix io.Reader implementation

Do not throw away the rest of Int63 value used for
generation random bytes. Save it in Rand struct and
re-use during the next Read call.

Fixes #16124

Change-Id: Ic70bd80c3c3a6590e60ac615e8b3c2324589bea3
Reviewed-on: https://go-review.googlesource.com/24251
Reviewed-by: Keith Randall <khr@golang.org>
Reviewed-by: Joe Tsai <thebrokentoaster@gmail.com>
Run-TryBot: Joe Tsai <thebrokentoaster@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>

3 files changed, 84 insertions, 28 deletions

diff --git a/src/math/rand/rand.go b/src/math/rand/rand.go
index 05a56e0f04..8f31b0ea9d 100644
--- a/src/math/rand/rand.go
+++ b/src/math/rand/rand.go
@@ -33,14 +33,26 @@ func NewSource(seed int64) Source {
 // A Rand is a source of random numbers.
 type Rand struct {
 	src Source
+
+	// readVal contains remainder of 63-bit integer used for bytes
+	// generation during most recent Read call.
+	// It is saved so next Read call can start where the previous
+	// one finished.
+	readVal int64
+	// readPos indicates the number of low-order bytes of readVal
+	// that are still valid.
+	readPos int8
 }
 
 // New returns a new Rand that uses random values from src
 // to generate other random values.
-func New(src Source) *Rand { return &Rand{src} }
+func New(src Source) *Rand { return &Rand{src: src} }
 
 // Seed uses the provided seed value to initialize the generator to a deterministic state.
-func (r *Rand) Seed(seed int64) { r.src.Seed(seed) }
+func (r *Rand) Seed(seed int64) {
+	r.src.Seed(seed)
+	r.readPos = 0
+}
 
 // Int63 returns a non-negative pseudo-random 63-bit integer as an int64.
 func (r *Rand) Int63() int64 { return r.src.Int63() }
@@ -161,14 +173,20 @@ func (r *Rand) Perm(n int) []int {
 // Read generates len(p) random bytes and writes them into p. It
 // always returns len(p) and a nil error.
 func (r *Rand) Read(p []byte) (n int, err error) {
-	for i := 0; i < len(p); i += 7 {
-		val := r.src.Int63()
-		for j := 0; i+j < len(p) && j < 7; j++ {
-			p[i+j] = byte(val)
-			val >>= 8
+	pos := r.readPos
+	val := r.readVal
+	for n = 0; n < len(p); n++ {
+		if pos == 0 {
+			val = r.Int63()
+			pos = 7
 		}
+		p[n] = byte(val)
+		val >>= 8
+		pos--
 	}
-	return len(p), nil
+	r.readPos = pos
+	r.readVal = val
+	return
 }
 
 /*
diff --git a/src/math/rand/rand_test.go b/src/math/rand/rand_test.go
index 8d68335fdd..6f312792d0 100644
--- a/src/math/rand/rand_test.go
+++ b/src/math/rand/rand_test.go
@@ -5,13 +5,16 @@
 package rand
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"internal/testenv"
+	"io"
 	"math"
 	"os"
 	"runtime"
 	"testing"
+	"testing/iotest"
 )
 
 const (
@@ -373,7 +376,7 @@ func testReadUniformity(t *testing.T, n int, seed int64) {
 	checkSampleDistribution(t, samples, expected)
 }
 
-func TestRead(t *testing.T) {
+func TestReadUniformity(t *testing.T) {
 	testBufferSizes := []int{
 		2, 4, 7, 64, 1024, 1 << 16, 1 << 20,
 	}
@@ -394,7 +397,42 @@ func TestReadEmpty(t *testing.T) {
 	if n != 0 {
 		t.Errorf("Read into empty buffer returned unexpected n of %d", n)
 	}
+}
+
+func TestReadByOneByte(t *testing.T) {
+	r := New(NewSource(1))
+	b1 := make([]byte, 100)
+	_, err := io.ReadFull(iotest.OneByteReader(r), b1)
+	if err != nil {
+		t.Errorf("read by one byte: %v", err)
+	}
+	r = New(NewSource(1))
+	b2 := make([]byte, 100)
+	_, err = r.Read(b2)
+	if err != nil {
+		t.Errorf("read: %v", err)
+	}
+	if !bytes.Equal(b1, b2) {
+		t.Errorf("read by one byte vs single read:\n%x\n%x", b1, b2)
+	}
+}
 
+func TestReadSeedReset(t *testing.T) {
+	r := New(NewSource(42))
+	b1 := make([]byte, 128)
+	_, err := r.Read(b1)
+	if err != nil {
+		t.Errorf("read: %v", err)
+	}
+	r.Seed(42)
+	b2 := make([]byte, 128)
+	_, err = r.Read(b2)
+	if err != nil {
+		t.Errorf("read: %v", err)
+	}
+	if !bytes.Equal(b1, b2) {
+		t.Errorf("mismatch after re-seed:\n%x\n%x", b1, b2)
+	}
 }
 
 // Benchmarks
diff --git a/src/math/rand/regress_test.go b/src/math/rand/regress_test.go
index fe9c4beff7..4dd965ce1e 100644
--- a/src/math/rand/regress_test.go
+++ b/src/math/rand/regress_test.go
@@ -342,25 +342,25 @@ var regressGolden = []interface{}{
 	[]int{8, 7, 5, 3, 4, 6, 0, 1, 2},    // Perm(9)
 	[]int{1, 0, 2, 5, 7, 6, 9, 8, 3, 4}, // Perm(10)
 	[]byte{0x1},                         // Read([0])
-	[]byte{0xc0, 0x41, 0xd3, 0xff, 0x12, 0x4, 0x5b},                   // Read([0 0 0 0 0 0 0])
-	[]byte{0x73, 0xc8, 0x6e, 0x4f, 0xf9, 0x5f, 0xf6, 0x62},            // Read([0 0 0 0 0 0 0 0])
-	[]byte{0x4a, 0x2d, 0xb, 0x75, 0xfb, 0x18, 0xd, 0xaf, 0x48},        // Read([0 0 0 0 0 0 0 0 0])
-	[]byte{0x39, 0x46, 0x51, 0x85, 0xf, 0xd4, 0xa1, 0x78, 0x89, 0x2e}, // Read([0 0 0 0 0 0 0 0 0 0])
-	[]byte{0x51}, // Read([0])
-	[]byte{0x4e, 0xe2, 0xd3, 0xd0, 0xd0, 0xde, 0x6b},                   // Read([0 0 0 0 0 0 0])
-	[]byte{0xf8, 0xf9, 0xb4, 0x4c, 0xe8, 0x5f, 0xf0, 0x44},             // Read([0 0 0 0 0 0 0 0])
-	[]byte{0x3b, 0xbf, 0x85, 0x7a, 0xab, 0x99, 0xc5, 0xb2, 0x52},       // Read([0 0 0 0 0 0 0 0 0])
-	[]byte{0xa8, 0xae, 0xb7, 0x9e, 0xf8, 0x56, 0xf6, 0x59, 0xc1, 0x8f}, // Read([0 0 0 0 0 0 0 0 0 0])
-	[]byte{0xc7}, // Read([0])
-	[]byte{0x5f, 0x67, 0xcf, 0xe2, 0x42, 0xcf, 0x3c},                   // Read([0 0 0 0 0 0 0])
-	[]byte{0xc3, 0x54, 0xf3, 0xed, 0xe2, 0xd6, 0xbe, 0xcc},             // Read([0 0 0 0 0 0 0 0])
-	[]byte{0x6a, 0x9f, 0x4a, 0x57, 0x8b, 0xcb, 0x9e, 0xf2, 0xd4},       // Read([0 0 0 0 0 0 0 0 0])
-	[]byte{0x6d, 0x29, 0x97, 0x61, 0xea, 0x9e, 0x4f, 0x5a, 0xa6, 0xae}, // Read([0 0 0 0 0 0 0 0 0 0])
-	[]byte{0xaa}, // Read([0])
-	[]byte{0x20, 0xef, 0xcd, 0x6c, 0xea, 0x84, 0xb6},                   // Read([0 0 0 0 0 0 0])
-	[]byte{0x92, 0x5e, 0x60, 0x7b, 0xe0, 0x63, 0x71, 0x6f},             // Read([0 0 0 0 0 0 0 0])
-	[]byte{0x4, 0x5c, 0x3f, 0x0, 0xf, 0x8a, 0x79, 0x6b, 0xce},          // Read([0 0 0 0 0 0 0 0 0])
-	[]byte{0xaa, 0xca, 0xee, 0xdf, 0xad, 0x5b, 0x50, 0x66, 0x64, 0xe8}, // Read([0 0 0 0 0 0 0 0 0 0])
+	[]byte{0x94, 0xfd, 0xc2, 0xfa, 0x2f, 0xfc, 0xc0},                 // Read([0 0 0 0 0 0 0])
+	[]byte{0x41, 0xd3, 0xff, 0x12, 0x4, 0x5b, 0x73, 0xc8},            // Read([0 0 0 0 0 0 0 0])
+	[]byte{0x6e, 0x4f, 0xf9, 0x5f, 0xf6, 0x62, 0xa5, 0xee, 0xe8},     // Read([0 0 0 0 0 0 0 0 0])
+	[]byte{0x2a, 0xbd, 0xf4, 0x4a, 0x2d, 0xb, 0x75, 0xfb, 0x18, 0xd}, // Read([0 0 0 0 0 0 0 0 0 0])
+	[]byte{0xaf},                                                      // Read([0])
+	[]byte{0x48, 0xa7, 0x9e, 0xe0, 0xb1, 0xd, 0x39},                   // Read([0 0 0 0 0 0 0])
+	[]byte{0x46, 0x51, 0x85, 0xf, 0xd4, 0xa1, 0x78, 0x89},             // Read([0 0 0 0 0 0 0 0])
+	[]byte{0x2e, 0xe2, 0x85, 0xec, 0xe1, 0x51, 0x14, 0x55, 0x78},      // Read([0 0 0 0 0 0 0 0 0])
+	[]byte{0x8, 0x75, 0xd6, 0x4e, 0xe2, 0xd3, 0xd0, 0xd0, 0xde, 0x6b}, // Read([0 0 0 0 0 0 0 0 0 0])
+	[]byte{0xf8}, // Read([0])
+	[]byte{0xf9, 0xb4, 0x4c, 0xe8, 0x5f, 0xf0, 0x44},                   // Read([0 0 0 0 0 0 0])
+	[]byte{0xc6, 0xb1, 0xf8, 0x3b, 0x8e, 0x88, 0x3b, 0xbf},             // Read([0 0 0 0 0 0 0 0])
+	[]byte{0x85, 0x7a, 0xab, 0x99, 0xc5, 0xb2, 0x52, 0xc7, 0x42},       // Read([0 0 0 0 0 0 0 0 0])
+	[]byte{0x9c, 0x32, 0xf3, 0xa8, 0xae, 0xb7, 0x9e, 0xf8, 0x56, 0xf6}, // Read([0 0 0 0 0 0 0 0 0 0])
+	[]byte{0x59},                                                       // Read([0])
+	[]byte{0xc1, 0x8f, 0xd, 0xce, 0xcc, 0x77, 0xc7},                    // Read([0 0 0 0 0 0 0])
+	[]byte{0x5e, 0x7a, 0x81, 0xbf, 0xde, 0x27, 0x5f, 0x67},             // Read([0 0 0 0 0 0 0 0])
+	[]byte{0xcf, 0xe2, 0x42, 0xcf, 0x3c, 0xc3, 0x54, 0xf3, 0xed},       // Read([0 0 0 0 0 0 0 0 0])
+	[]byte{0xe2, 0xd6, 0xbe, 0xcc, 0x4e, 0xa3, 0xae, 0x5e, 0x88, 0x52}, // Read([0 0 0 0 0 0 0 0 0 0])
 	uint32(4059586549),                                                 // Uint32()
 	uint32(1052117029),                                                 // Uint32()
 	uint32(2817310706),                                                 // Uint32()