From a526ec1569984fd452acb09c8c48c9d972a442d3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Wed, 15 Jun 2022 21:49:50 +0100
Subject: [PATCH] go/printer: avoid copying bytes when counting node sizes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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 | 36 +++++++++++++++++++++++++-----------
 1 file changed, 25 insertions(+), 11 deletions(-)

diff --git a/src/go/printer/nodes.go b/src/go/printer/nodes.go
index c7cab46f17f..c167b5f137a 100644
--- 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