bufio: Implement io.ReaderFrom for (*Writer).

This is part 2 of 2 for issue 4028.

benchmark                        old ns/op    new ns/op    delta
BenchmarkWriterCopyOptimal           53293        28326  -46.85%
BenchmarkWriterCopyUnoptimal         53757        30537  -43.19%
BenchmarkWriterCopyNoReadFrom        53192        36642  -31.11%

Fixes #4028.

R=nigeltao
CC=golang-dev
https://golang.org/cl/6565056

diff --git a/src/pkg/bufio/bufio.go b/src/pkg/bufio/bufio.go
index d6ba485fa3e88aeb95731f430891f8ae12a74bb1..d1c5a13bcab4758af03de7e7eb9f07d871541303 100644 (file)
--- a/src/pkg/bufio/bufio.go
+++ b/src/pkg/bufio/bufio.go
@@ -567,6 +567,35 @@ func (b *Writer) WriteString(s string) (int, error) {
        return nn, nil
 }
 
+// ReadFrom implements io.ReaderFrom.
+func (b *Writer) ReadFrom(r io.Reader) (n int64, err error) {
+       if err = b.Flush(); err != nil {
+               return 0, err
+       }
+       if w, ok := b.wr.(io.ReaderFrom); ok {
+               return w.ReadFrom(r)
+       }
+       var m int
+       for {
+               m, err = r.Read(b.buf[b.n:])
+               if m == 0 {
+                       break
+               }
+               b.n += m
+               n += int64(m)
+               if err1 := b.Flush(); err1 != nil {
+                       return n, err1
+               }
+               if err != nil {
+                       break
+               }
+       }
+       if err == io.EOF {
+               err = nil
+       }
+       return n, err
+}
+
 // buffered input and output
 
 // ReadWriter stores pointers to a Reader and a Writer.

diff --git a/src/pkg/bufio/bufio_test.go b/src/pkg/bufio/bufio_test.go
index 4e10207efbd4948a4e8ecd28972a073f41fbfc34..3d07639e2aa704fbd5e73fd40f967029d2f39fab 100644 (file)
--- a/src/pkg/bufio/bufio_test.go
+++ b/src/pkg/bufio/bufio_test.go
@@ -763,8 +763,8 @@ func testReadLineNewlines(t *testing.T, input string, expect []readLineResult) {
        }
 }
 
-func TestReaderWriteTo(t *testing.T) {
-       input := make([]byte, 8192)
+func createTestInput(n int) []byte {
+       input := make([]byte, n)
        for i := range input {
                // 101 and 251 are arbitrary prime numbers.
                // The idea is to create an input sequence
@@ -774,7 +774,12 @@ func TestReaderWriteTo(t *testing.T) {
                        input[i] ^= byte(i / 101)
                }
        }
-       r := NewReader(bytes.NewBuffer(input))
+       return input
+}
+
+func TestReaderWriteTo(t *testing.T) {
+       input := createTestInput(8192)
+       r := NewReader(&onlyReader{bytes.NewBuffer(input)})
        w := new(bytes.Buffer)
        if n, err := r.WriteTo(w); err != nil || n != int64(len(input)) {
                t.Fatalf("r.WriteTo(w) = %d, %v, want %d, nil", n, err, len(input))
@@ -817,6 +822,65 @@ func TestReaderWriteToErrors(t *testing.T) {
        }
 }
 
+func TestWriterReadFrom(t *testing.T) {
+       ws := []func(io.Writer) io.Writer{
+               func(w io.Writer) io.Writer { return &onlyWriter{w} },
+               func(w io.Writer) io.Writer { return w },
+       }
+
+       rs := []func(io.Reader) io.Reader{
+               iotest.DataErrReader,
+               func(r io.Reader) io.Reader { return r },
+       }
+
+       for ri, rfunc := range rs {
+               for wi, wfunc := range ws {
+                       input := createTestInput(8192)
+                       b := new(bytes.Buffer)
+                       w := NewWriter(wfunc(b))
+                       r := rfunc(bytes.NewBuffer(input))
+                       if n, err := w.ReadFrom(r); err != nil || n != int64(len(input)) {
+                               t.Errorf("ws[%d],rs[%d]: w.ReadFrom(r) = %d, %v, want %d, nil", wi, ri, n, err, len(input))
+                               continue
+                       }
+                       if got, want := b.String(), string(input); got != want {
+                               t.Errorf("ws[%d], rs[%d]:\ngot  %q\nwant %q\n", wi, ri, got, want)
+                       }
+               }
+       }
+}
+
+type errorReaderFromTest struct {
+       rn, wn     int
+       rerr, werr error
+       expected   error
+}
+
+func (r errorReaderFromTest) Read(p []byte) (int, error) {
+       return len(p) * r.rn, r.rerr
+}
+
+func (w errorReaderFromTest) Write(p []byte) (int, error) {
+       return len(p) * w.wn, w.werr
+}
+
+var errorReaderFromTests = []errorReaderFromTest{
+       {0, 1, io.EOF, nil, nil},
+       {1, 1, io.EOF, nil, nil},
+       {0, 1, io.ErrClosedPipe, nil, io.ErrClosedPipe},
+       {0, 0, io.ErrClosedPipe, io.ErrShortWrite, io.ErrClosedPipe},
+       {1, 0, nil, io.ErrShortWrite, io.ErrShortWrite},
+}
+
+func TestWriterReadFromErrors(t *testing.T) {
+       for i, rw := range errorReaderFromTests {
+               w := NewWriter(rw)
+               if _, err := w.ReadFrom(rw); err != rw.expected {
+                       t.Errorf("w.ReadFrom(errorReaderFromTests[%d]) = _, %v, want _,%v", i, err, rw.expected)
+               }
+       }
+}
+
 // An onlyReader only implements io.Reader, no matter what other methods the underlying implementation may have.
 type onlyReader struct {
        r io.Reader
@@ -866,3 +930,34 @@ func BenchmarkReaderCopyNoWriteTo(b *testing.B) {
                io.Copy(dst, src)
        }
 }
+
+func BenchmarkWriterCopyOptimal(b *testing.B) {
+       // Optimal case is where the underlying writer implements io.ReaderFrom
+       for i := 0; i < b.N; i++ {
+               b.StopTimer()
+               src := &onlyReader{bytes.NewBuffer(make([]byte, 8192))}
+               dst := NewWriter(new(bytes.Buffer))
+               b.StartTimer()
+               io.Copy(dst, src)
+       }
+}
+
+func BenchmarkWriterCopyUnoptimal(b *testing.B) {
+       for i := 0; i < b.N; i++ {
+               b.StopTimer()
+               src := &onlyReader{bytes.NewBuffer(make([]byte, 8192))}
+               dst := NewWriter(&onlyWriter{new(bytes.Buffer)})
+               b.StartTimer()
+               io.Copy(dst, src)
+       }
+}
+
+func BenchmarkWriterCopyNoReadFrom(b *testing.B) {
+       for i := 0; i < b.N; i++ {
+               b.StopTimer()
+               src := &onlyReader{bytes.NewBuffer(make([]byte, 8192))}
+               dst := &onlyWriter{NewWriter(new(bytes.Buffer))}
+               b.StartTimer()
+               io.Copy(dst, src)
+       }
+}