go/printer: avoid copying bytes when counting node sizes

When making the recursive call to fprint,
we just need to know how many bytes were written
and whether multiple lines were written.

We don't need a buffer to accomplish those;
a custom writer can keep track of the two in a cheap way,
avoiding the allocation of a byte slice and copying bytes.

	name      old time/op    new time/op    delta
	Print-16    6.28ms ± 2%    6.12ms ± 1%  -2.50%  (p=0.000 n=10+9)

	name      old speed      new speed      delta
	Print-16  8.26MB/s ± 3%  8.47MB/s ± 1%  +2.56%  (p=0.000 n=10+9)

	name      old alloc/op   new alloc/op   delta
	Print-16     483kB ± 0%     443kB ± 0%  -8.20%  (p=0.000 n=10+10)

	name      old allocs/op  new allocs/op  delta
	Print-16     17.8k ± 0%     17.3k ± 0%  -2.31%  (p=0.000 n=9+10)

Change-Id: Ib8411ae6738a2acae6af6d185da71727ce2eb97a
Reviewed-on: https://go-review.googlesource.com/c/go/+/412555
Reviewed-by: Robert Griesemer <gri@google.com>
Reviewed-by: Than McIntosh <thanm@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Run-TryBot: Daniel Martí <mvdan@mvdan.cc>

src/go/printer/nodes.go

@@ -9,7 +9,6 @@
 package printer
 
 import (
-	"bytes"
 	"go/ast"
 	"go/token"
 	"math"
@@ -1720,6 +1719,26 @@ func (p *printer) genDecl(d *ast.GenDecl) {
 	}
 }
 
+// sizeCounter is an io.Writer which counts the number of bytes written,
+// as well as whether a newline character was seen.
+type sizeCounter struct {
+	hasNewline bool
+	size       int
+}
+
+func (c *sizeCounter) Write(p []byte) (int, error) {
+	if !c.hasNewline {
+		for _, b := range p {
+			if b == '\n' || b == '\f' {
+				c.hasNewline = true
+				break
+			}
+		}
+	}
+	c.size += len(p)
+	return len(p), nil
+}
+
 // nodeSize determines the size of n in chars after formatting.
 // The result is <= maxSize if the node fits on one line with at
 // most maxSize chars and the formatted output doesn't contain
@@ -1740,18 +1759,13 @@ func (p *printer) nodeSize(n ast.Node, maxSize int) (size int) {
 	// style so that we always get the same decision; print
 	// in RawFormat
 	cfg := Config{Mode: RawFormat}
-	var buf bytes.Buffer
+	var counter sizeCounter
-	if err := cfg.fprint(&buf, p.fset, n, p.nodeSizes); err != nil {
+	if err := cfg.fprint(&counter, p.fset, n, p.nodeSizes); err != nil {
 		return
 	}
-	if buf.Len() <= maxSize {
+	if counter.size <= maxSize && !counter.hasNewline {
-		for _, ch := range buf.Bytes() {
+		// n fits in a single line
-			switch ch {
+		size = counter.size
-			case '\n', '\f':
-				return // does not fit in a single line
-			}
-		}
-		size = buf.Len() // n fits
 		p.nodeSizes[n] = size
 	}
 	return