bufio: add Writer.AvailableBuffer

This adds a new Writer.AvailableBuffer method that returns
an empty buffer with a possibly non-empty capacity for use
with append-like APIs.

The typical usage pattern is something like:
	b := bw.AvailableBuffer()
	b = appendValue(b, v)
	bw.Write(b)

It allows logic combining append-like APIs with bufio.Writer to avoid
needing to allocate and manage buffers themselves and allows the
append-like APIs to directly write into the buffer for a bufio.Writer.

Fixes #47527

Change-Id: I9cd169f3f8e8c7cd40818caf3daf1944c826fc66
Reviewed-on: https://go-review.googlesource.com/c/go/+/345569
Trust: Joe Tsai <joetsai@digital-static.net>
Run-TryBot: Joe Tsai <joetsai@digital-static.net>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>

diff --git a/src/bufio/bufio.go b/src/bufio/bufio.go
index 391ecf46b339e23df728d659d63127734a7fce1a..506b84f6baf94f9d4bd82c81d84cc492ae64d3c6 100644 (file)
--- a/src/bufio/bufio.go
+++ b/src/bufio/bufio.go
@@ -633,6 +633,14 @@ func (b *Writer) Flush() error {
 // Available returns how many bytes are unused in the buffer.
 func (b *Writer) Available() int { return len(b.buf) - b.n }
 
+// AvailableBuffer returns an empty buffer with b.Available() capacity.
+// This buffer is intended to be appended to and
+// passed to an immediately succeeding Write call.
+// The buffer is only valid until the next write operation on b.
+func (b *Writer) AvailableBuffer() []byte {
+       return b.buf[b.n:][:0]
+}
+
 // Buffered returns the number of bytes that have been written into the current buffer.
 func (b *Writer) Buffered() int { return b.n }
 

diff --git a/src/bufio/bufio_test.go b/src/bufio/bufio_test.go
index eb5136c9ea858c7bf6b21a6eefc4bdcb1e9bba20..04a810c206500d846174c291e2b25fc4af59848f 100644 (file)
--- a/src/bufio/bufio_test.go
+++ b/src/bufio/bufio_test.go
@@ -10,6 +10,8 @@ import (
        "errors"
        "fmt"
        "io"
+       "math/rand"
+       "strconv"
        "strings"
        "testing"
        "testing/iotest"
@@ -608,6 +610,37 @@ func TestWriter(t *testing.T) {
        }
 }
 
+func TestWriterAppend(t *testing.T) {
+       got := new(bytes.Buffer)
+       var want []byte
+       rn := rand.New(rand.NewSource(0))
+       w := NewWriterSize(got, 64)
+       for i := 0; i < 100; i++ {
+               // Obtain a buffer to append to.
+               b := w.AvailableBuffer()
+               if w.Available() != cap(b) {
+                       t.Fatalf("Available() = %v, want %v", w.Available(), cap(b))
+               }
+
+               // While not recommended, it is valid to append to a shifted buffer.
+               // This forces Write to copy the the input.
+               if rn.Intn(8) == 0 && cap(b) > 0 {
+                       b = b[1:1:cap(b)]
+               }
+
+               // Append a random integer of varying width.
+               n := int64(rn.Intn(1 << rn.Intn(30)))
+               want = append(strconv.AppendInt(want, n, 10), ' ')
+               b = append(strconv.AppendInt(b, n, 10), ' ')
+               w.Write(b)
+       }
+       w.Flush()
+
+       if !bytes.Equal(got.Bytes(), want) {
+               t.Errorf("output mismatch:\ngot  %s\nwant %s", got.Bytes(), want)
+       }
+}
+
 // Check that write errors are returned properly.
 
 type errorWriterTest struct {

diff --git a/src/bufio/example_test.go b/src/bufio/example_test.go
index 8885d40549f1532010a320d00964f19bffbb5846..a864d11012e77f96bc2a24dc692ce967c73270e3 100644 (file)
--- a/src/bufio/example_test.go
+++ b/src/bufio/example_test.go
@@ -20,6 +20,18 @@ func ExampleWriter() {
        // Output: Hello, world!
 }
 
+func ExampleWriter_AvailableBuffer() {
+       w := bufio.NewWriter(os.Stdout)
+       for _, i := range []int64{1, 2, 3, 4} {
+               b := w.AvailableBuffer()
+               b = strconv.AppendInt(b, i, 10)
+               b = append(b, ' ')
+               w.Write(b)
+       }
+       w.Flush()
+       // Output: 1 2 3 4
+}
+
 // The simplest use of a Scanner, to read standard input as a set of lines.
 func ExampleScanner_lines() {
        scanner := bufio.NewScanner(os.Stdin)