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>

diff --git a/src/go/printer/nodes.go b/src/go/printer/nodes.go
index c7cab46f17fc9c5c28cea877d03b6a2833e1506e..c167b5f137a6b5eb24287822b701c624bd1ea48d 100644 (file)
--- a/src/go/printer/nodes.go
+++ b/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
-       if err := cfg.fprint(&buf, p.fset, n, p.nodeSizes); err != nil {
+       var counter sizeCounter
+       if err := cfg.fprint(&counter, p.fset, n, p.nodeSizes); err != nil {
                return
        }
-       if buf.Len() <= maxSize {
-               for _, ch := range buf.Bytes() {
-                       switch ch {
-                       case '\n', '\f':
-                               return // does not fit in a single line
-                       }
-               }
-               size = buf.Len() // n fits
+       if counter.size <= maxSize && !counter.hasNewline {
+               // n fits in a single line
+               size = counter.size
                p.nodeSizes[n] = size
        }
        return