go/printer: reduce allocations to improve performance

First, we know that Go source files almost always weigh at least a few
kilobytes, so we can kickstart the output buffer to be a reasonable size
and reduce the initial number of incremental allocations and copies when
appending bytes or strings to output.

Second, in nodeSize we use a nested printer, but we don't actually need
its printed bytes - we only need to know how many bytes it prints.
For that reason, use a throwaway buffer: the part of our output buffer
between length and capacity, as we haven't used it yet.

Third, use a sync.Pool to reuse allocated printers.
The current API doesn't allow reusing printers,
and some programs like gofmt will print many files in sequence.

Those changes combined result in a modest reduction in allocations and
CPU usage. The benchmark uses testdata/parser.go, which has just over
two thousand lines of code, which is pretty standard size-wise.

We also split the Print benchmark to cover both a medium-sized ast.File
as well as a pretty small ast.Decl node. The latter is a somewhat common
scenario in gopls, which has code actions which alter small bits of the
AST and print them back out to rewrite only a few lines in a file.

	name          old time/op    new time/op     delta
	PrintFile-16    5.43ms ± 1%     4.85ms ± 3%  -10.68%  (p=0.000 n=9+10)
	PrintDecl-16    19.1µs ± 0%     18.5µs ± 1%   -3.04%  (p=0.000 n=10+10)

	name          old speed      new speed       delta
	PrintFile-16  9.56MB/s ± 1%  10.69MB/s ± 3%  +11.81%  (p=0.000 n=8+10)
	PrintDecl-16  1.67MB/s ± 0%   1.73MB/s ± 1%   +3.05%  (p=0.000 n=10+10)

	name          old alloc/op   new alloc/op    delta
	PrintFile-16     332kB ± 0%      107kB ± 2%  -67.87%  (p=0.000 n=10+10)
	PrintDecl-16    3.92kB ± 0%     3.28kB ± 0%  -16.38%  (p=0.000 n=10+10)

	name          old allocs/op  new allocs/op   delta
	PrintFile-16     3.45k ± 0%      2.42k ± 0%  -29.90%  (p=0.000 n=10+10)
	PrintDecl-16      56.0 ± 0%       46.0 ± 0%  -17.86%  (p=0.000 n=10+10)

Change-Id: I475a3babca77532b2d51888f49710f74763d81d2
Reviewed-on: https://go-review.googlesource.com/c/go/+/424924
Run-TryBot: Daniel Martí <mvdan@mvdan.cc>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Robert Griesemer <gri@google.com>
Reviewed-by: Alan Donovan <adonovan@google.com>

src/go/printer/performance_test.go

@@ -11,6 +11,7 @@ import (
 	"bytes"
 	"go/ast"
 	"go/parser"
+	"go/token"
 	"io"
 	"log"
 	"os"
@@ -18,12 +19,15 @@ import (
 )
 
 var (
-	testfile *ast.File
-	testsize int64
+	fileNode *ast.File
+	fileSize int64
+
+	declNode ast.Decl
+	declSize int64
 )
 
-func testprint(out io.Writer, file *ast.File) {
-	if err := (&Config{TabIndent | UseSpaces | normalizeNumbers, 8, 0}).Fprint(out, fset, file); err != nil {
+func testprint(out io.Writer, node ast.Node) {
+	if err := (&Config{TabIndent | UseSpaces | normalizeNumbers, 8, 0}).Fprint(out, fset, node); err != nil {
 		log.Fatalf("print error: %s", err)
 	}
 }
@@ -48,17 +52,40 @@ func initialize() {
 		log.Fatalf("print error: %s not idempotent", filename)
 	}
 
-	testfile = file
-	testsize = int64(len(src))
+	fileNode = file
+	fileSize = int64(len(src))
+
+	for _, decl := range file.Decls {
+		// The first global variable, which is pretty short:
+		//
+		//	var unresolved = new(ast.Object)
+		if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.VAR {
+			declNode = decl
+			declSize = int64(fset.Position(decl.End()).Offset - fset.Position(decl.Pos()).Offset)
+			break
+		}
+
+	}
 }
 
-func BenchmarkPrint(b *testing.B) {
-	if testfile == nil {
+func BenchmarkPrintFile(b *testing.B) {
+	if fileNode == nil {
 		initialize()
 	}
 	b.ReportAllocs()
-	b.SetBytes(testsize)
+	b.SetBytes(fileSize)
 	for i := 0; i < b.N; i++ {
-		testprint(io.Discard, testfile)
+		testprint(io.Discard, fileNode)
+	}
+}
+
+func BenchmarkPrintDecl(b *testing.B) {
+	if declNode == nil {
+		initialize()
+	}
+	b.ReportAllocs()
+	b.SetBytes(declSize)
+	for i := 0; i < b.N; i++ {
+		testprint(io.Discard, declNode)
 	}
 }

src/go/printer/printer.go

@@ -13,6 +13,7 @@ import (
 	"io"
 	"os"
 	"strings"
+	"sync"
 	"text/tabwriter"
 	"unicode"
 )
@@ -94,16 +95,6 @@ type printer struct {
 	cachedLine int // line corresponding to cachedPos
 }
 
-func (p *printer) init(cfg *Config, fset *token.FileSet, nodeSizes map[ast.Node]int) {
-	p.Config = *cfg
-	p.fset = fset
-	p.pos = token.Position{Line: 1, Column: 1}
-	p.out = token.Position{Line: 1, Column: 1}
-	p.wsbuf = make([]whiteSpace, 0, 16) // whitespace sequences are short
-	p.nodeSizes = nodeSizes
-	p.cachedPos = -1
-}
-
 func (p *printer) internalError(msg ...any) {
 	if debug {
 		fmt.Print(p.pos.String() + ": ")
@@ -1324,11 +1315,47 @@ type Config struct {
 	Indent   int  // default: 0 (all code is indented at least by this much)
 }
 
+var printerPool = sync.Pool{
+	New: func() any {
+		return &printer{
+			// Whitespace sequences are short.
+			wsbuf: make([]whiteSpace, 0, 16),
+			// We start the printer with a 16K output buffer, which is currently
+			// larger than about 80% of Go files in the standard library.
+			output: make([]byte, 0, 16<<10),
+		}
+	},
+}
+
+func newPrinter(cfg *Config, fset *token.FileSet, nodeSizes map[ast.Node]int) *printer {
+	p := printerPool.Get().(*printer)
+	*p = printer{
+		Config:    *cfg,
+		fset:      fset,
+		pos:       token.Position{Line: 1, Column: 1},
+		out:       token.Position{Line: 1, Column: 1},
+		wsbuf:     p.wsbuf[:0],
+		nodeSizes: nodeSizes,
+		cachedPos: -1,
+		output:    p.output[:0],
+	}
+	return p
+}
+
+func (p *printer) free() {
+	// Hard limit on buffer size; see https://golang.org/issue/23199.
+	if cap(p.output) > 64<<10 {
+		return
+	}
+
+	printerPool.Put(p)
+}
+
 // fprint implements Fprint and takes a nodesSizes map for setting up the printer state.
 func (cfg *Config) fprint(output io.Writer, fset *token.FileSet, node any, nodeSizes map[ast.Node]int) (err error) {
 	// print node
-	var p printer
-	p.init(cfg, fset, nodeSizes)
+	p := newPrinter(cfg, fset, nodeSizes)
+	defer p.free()
 	if err = p.printNode(node); err != nil {
 		return
 	}