diff --git a/src/bufio/bufio.go b/src/bufio/bufio.go index f9b762be8d4966..ae72214718083f 100644 --- a/src/bufio/bufio.go +++ b/src/bufio/bufio.go @@ -745,20 +745,36 @@ func (b *Writer) WriteRune(r rune) (size int, err error) { // If the count is less than len(s), it also returns an error explaining // why the write is short. func (b *Writer) WriteString(s string) (int, error) { - var sw io.StringWriter - tryStringWriter := true + if b.err != nil { + return 0, b.err + } + + if len(s) <= b.Available() { + // Fast path: the whole string fits in the buffer. + n := copy(b.buf[b.n:], s) + b.n += n + return n, nil + } + + sw, ok := b.wr.(io.StringWriter) nn := 0 - for len(s) > b.Available() && b.err == nil { + if ok && b.Buffered() == 0 { + // Large write, empty buffer, and the underlying writer supports + // WriteString: forward the write to the underlying StringWriter. + // This avoids an extra copy. + var n int + n, b.err = sw.WriteString(s) + if n == len(s) || b.err != nil { + return n, b.err + } + nn = n + s = s[n:] + } + + for { var n int - if b.Buffered() == 0 && sw == nil && tryStringWriter { - // Check at most once whether b.wr is a StringWriter. - sw, tryStringWriter = b.wr.(io.StringWriter) - } - if b.Buffered() == 0 && tryStringWriter { - // Large write, empty buffer, and the underlying writer supports - // WriteString: forward the write to the underlying StringWriter. - // This avoids an extra copy. + if b.Buffered() == 0 && ok { n, b.err = sw.WriteString(s) } else { n = copy(b.buf[b.n:], s) @@ -767,6 +783,9 @@ func (b *Writer) WriteString(s string) (int, error) { } nn += n s = s[n:] + if len(s) <= b.Available() || b.err != nil { + break + } } if b.err != nil { return nn, b.err diff --git a/src/bufio/bufio_test.go b/src/bufio/bufio_test.go index 742e1954256903..7e11fb3b81b31c 100644 --- a/src/bufio/bufio_test.go +++ b/src/bufio/bufio_test.go @@ -1997,3 +1997,20 @@ func BenchmarkWriterFlush(b *testing.B) { bw.Flush() } } + +func BenchmarkWriteString(b *testing.B) { + small := strings.Repeat("x", 50) + huge := strings.Repeat("x", 8<<10) + for _, tc := range []struct{ name, s string }{ + {"small", small}, + {"huge", huge}, + } { + b.Run(tc.name, func(b *testing.B) { + b.ReportAllocs() + bw := NewWriter(io.Discard) + for b.Loop() { + bw.WriteString(tc.s) + } + }) + } +}