bufio: add fast path to WriteString for strings that fit in the buffer

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

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)
+			}
+		})
+	}
+}