math/rand/v2: add Read method to ChaCha8 source

api/next/67059.txt

@@ -0,0 +1 @@
+pkg math/rand/v2, method (*ChaCha8) Read([]uint8) (int, error) #67059

doc/next/6-stdlib/99-minor/math/rand/v2/67059.md

@@ -0,0 +1 @@
+[ChaCha8] source implements the [io.Reader] interface.

src/internal/chacha8rand/chacha8.go

@@ -7,7 +7,11 @@
 // and must have minimal dependencies.
 package chacha8rand
 
-import "internal/byteorder"
+import (
+	"internal/byteorder"
+	"internal/goarch"
+	"unsafe"
+)
 
 const (
 	ctrInc = 4  // increment counter by 4 between block calls
@@ -158,3 +162,40 @@ func Unmarshal(s *State, data []byte) error {
 	}
 	return nil
 }
+
+func (s *State) FillRand(b []byte) {
+	if s.i == s.n {
+		s.Refill()
+	}
+
+	for {
+		curRand := s.buf[s.i:s.n]
+
+		// Make sure that on little and big endian systems the bytes are filled in the same order.
+		if goarch.BigEndian {
+			for _, v := range curRand {
+				var tmp [8]byte
+				byteorder.LePutUint64(tmp[:], v)
+				n := copy(b, tmp[:])
+				s.i++
+				b = b[n:]
+
+				if len(b) == 0 {
+					return
+				}
+			}
+			s.Refill()
+			continue
+		}
+
+		curRandBytes := unsafe.Slice((*byte)(unsafe.Pointer(&curRand[0])), len(curRand)*8)
+		n := copy(b, curRandBytes)
+		s.i += (uint32(n) + 7) / 8
+		b = b[n:]
+
+		if len(b) == 0 {
+			return
+		}
+		s.Refill()
+	}
+}

src/internal/chacha8rand/rand_test.go

@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"fmt"
+	"internal/byteorder"
 	. "internal/chacha8rand"
 	"slices"
 	"testing"
@@ -30,6 +31,38 @@ func TestOutput(t *testing.T) {
 	}
 }
 
+func TestFillRand(t *testing.T) {
+	expect := make([]byte, 0, len(output)*8)
+	for _, v := range output {
+		expect = byteorder.LeAppendUint64(expect, v)
+	}
+
+	cases := []struct {
+		expect [][]byte
+	}{
+		{[][]byte{expect}},
+		{[][]byte{expect[:1], expect[8:9]}},
+		{[][]byte{expect[:4], expect[8:12]}},
+		{[][]byte{expect[:4], expect[8:12]}},
+		{[][]byte{expect[:8], expect[8:16]}},
+		{[][]byte{expect[:128], expect[128:256]}},
+		{[][]byte{expect[:128], expect[128:]}},
+		{[][]byte{expect[:100], expect[104:]}},
+	}
+
+	var s State
+	for _, tt := range cases {
+		s.Init(seed)
+		for _, expect := range tt.expect {
+			got := make([]byte, len(expect))
+			s.FillRand(got)
+			if !bytes.Equal(expect, got) {
+				t.Errorf("got = %#x; want = %#x", got, expect)
+			}
+		}
+	}
+}
+
 func TestMarshal(t *testing.T) {
 	var s State
 	s.Init(seed)

src/math/rand/v2/chacha8.go

@@ -44,3 +44,10 @@ func (c *ChaCha8) UnmarshalBinary(data []byte) error {
 func (c *ChaCha8) MarshalBinary() ([]byte, error) {
 	return chacha8rand.Marshal(&c.state), nil
 }
+
+// Read generates len(p) random bytes and writes them into p.
+// It never returns an error.
+func (c *ChaCha8) Read(p []byte) (n int, err error) {
+	c.state.FillRand(p)
+	return len(p), nil
+}

src/math/rand/v2/chacha8_test.go

@@ -6,6 +6,7 @@ package rand_test
 
 import (
 	. "math/rand/v2"
+	"strconv"
 	"testing"
 )
 
@@ -529,3 +530,30 @@ var chacha8marshal = []string{
 	"chacha8:\x00\x00\x00\x00\x00\x00\x00zK3\x9bB!,\x94\x9d\x975\xce'O_t\xee|\xb21\x87\xbb\xbb\xfd)\x8f\xe52\x01\vP\fk",
 	"chacha8:\x00\x00\x00\x00\x00\x00\x00{K3\x9bB!,\x94\x9d\x975\xce'O_t\xee|\xb21\x87\xbb\xbb\xfd)\x8f\xe52\x01\vP\fk",
 }
+
+func TestChaCha8Read(t *testing.T) {
+	c := NewChaCha8(chacha8seed)
+	n, err := c.Read(make([]byte, 128))
+	if n != 128 || err != nil {
+		t.Errorf("(*ChaCha8).Read(buf) = (%v, %v); want = (128, nil)", n, err)
+	}
+
+	n, err = c.Read(make([]byte, 0))
+	if n != 0 || err != nil {
+		t.Errorf("(*ChaCha8).Read(buf) = (%v, %v); want = (0, nil)", n, err)
+	}
+}
+
+func BenchmarkRead(b *testing.B) {
+	for _, v := range []int{1, 2, 3, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 1024 * 2, 1024 * 16} {
+		b.Run(strconv.FormatInt(int64(v), 10), func(b *testing.B) {
+			r := NewChaCha8(chacha8seed)
+			buf := make([]byte, v)
+			b.SetBytes(int64(v))
+			b.ResetTimer()
+			for n := b.N; n > 0; n-- {
+				r.Read(buf)
+			}
+		})
+	}
+}