cmd/compile: add astdump debug flag

This was extraordinarily useful for inlining work.
I have cleaned it up somewhat, and did some additional tweaks
after working on changes to bloop.

-gcflags=-d=astdump=SomeFunc
-gcflags=-d=astdump=SomeSubPkg.SomeFunc
-gcflags=-d=astdump=Some/Pkg.SomeFunc
-gcflags=-d=astdump=~YourRegExpHere

Change-Id: I3f98601ca96c87d6b191d4b64b264cd236e6d8bf
Reviewed-on: https://go-review.googlesource.com/c/go/+/629775
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Keith Randall <khr@google.com>
Reviewed-by: Keith Randall <khr@golang.org>

src/cmd/compile/internal/base/debug.go

@@ -18,6 +18,7 @@ var Debug DebugFlags
 type DebugFlags struct {
 	AlignHot              int    `help:"enable hot block alignment (currently requires -pgo)" concurrent:"ok"`
 	Append                int    `help:"print information about append compilation"`
+	AstDump               string `help:"for specified function/method, dump AST/IR at interesting points in compilation, to file pkg.func.ast. Use leading ~ for regular expression match."`
 	Checkptr              int    `help:"instrument unsafe pointer conversions\n0: instrumentation disabled\n1: conversions involving unsafe.Pointer are instrumented\n2: conversions to unsafe.Pointer force heap allocation" concurrent:"ok"`
 	Closure               int    `help:"print information about closure compilation"`
 	CompressInstructions  int    `help:"use compressed instructions when possible (if supported by architecture)"`

src/cmd/compile/internal/bloop/bloop.go

@@ -320,5 +320,9 @@ func Walk(pkg *ir.Package) {
 	for _, fn := range pkg.Funcs {
 		e := editor{false, fn}
 		ir.EditChildren(fn, e.edit)
+		if ir.MatchAstDump(fn, "bloop") {
+			ir.AstDump(fn, "bloop, "+ir.FuncName(fn))
+		}
 	}
+
 }

src/cmd/compile/internal/deadlocals/deadlocals.go

@@ -23,7 +23,11 @@ func Funcs(fns []*ir.Func) {
 	zero := ir.NewBasicLit(base.AutogeneratedPos, types.Types[types.TINT], constant.MakeInt64(0))
 
 	for _, fn := range fns {
+
 		if fn.IsClosure() {
+			if ir.MatchAstDump(fn, "deadlocals closure") {
+				ir.AstDump(fn, "deadlocals closure skipped, "+ir.FuncName(fn))
+			}
 			continue
 		}
 
@@ -50,6 +54,9 @@ func Funcs(fns []*ir.Func) {
 				k.Defn = nil
 			}
 		}
+		if ir.MatchAstDump(fn, "deadlocals") {
+			ir.AstDump(fn, "deadLocals, "+ir.FuncName(fn))
+		}
 	}
 }
 

src/cmd/compile/internal/escape/escape.go

@@ -380,6 +380,11 @@ func (b *batch) finish(fns []*ir.Func) {
 		}
 	}
 
+	for _, fn := range fns {
+		if ir.MatchAstDump(fn, "escape") {
+			ir.AstDump(fn, "escape, "+ir.FuncName(fn))
+		}
+	}
 }
 
 // inMutualBatch reports whether function fn is in the batch of

src/cmd/compile/internal/gc/compile.go

@@ -121,6 +121,9 @@ func prepareFunc(fn *ir.Func) {
 
 	ir.CurFunc = fn
 	walk.Walk(fn)
+	if ir.MatchAstDump(fn, "walk") {
+		ir.AstDump(fn, "walk, "+ir.FuncName(fn))
+	}
 	ir.CurFunc = nil // enforce no further uses of CurFunc
 
 	base.Ctxt.DwTextCount++

src/cmd/compile/internal/inline/inl.go

@@ -286,6 +286,8 @@ func CanInline(fn *ir.Func, profile *pgoir.Profile) {
 	// locals, and we use this map to produce a pruned Inline.Dcl
 	// list. See issue 25459 for more context.
 
+	dbg := ir.MatchAstDump(fn, "inline")
+
 	visitor := hairyVisitor{
 		curFunc:       fn,
 		debug:         isDebugFn(fn),
@@ -294,10 +296,17 @@ func CanInline(fn *ir.Func, profile *pgoir.Profile) {
 		maxBudget:     budget,
 		extraCallCost: cc,
 		profile:       profile,
+		dbg:           dbg, // Useful for downstream debugging
 	}
+
 	if visitor.tooHairy(fn) {
 		reason = visitor.reason
+		if dbg {
+			ir.AstDump(fn, "inline, too hairy because "+visitor.reason+", "+ir.FuncName(fn))
+		}
 		return
+	} else if dbg {
+		ir.AstDump(fn, "inline, OK, "+ir.FuncName(fn))
 	}
 
 	n.Func.Inl = &ir.Inline{
@@ -441,6 +450,7 @@ type hairyVisitor struct {
 	usedLocals    ir.NameSet
 	do            func(ir.Node) bool
 	profile       *pgoir.Profile
+	dbg           bool
 }
 
 func isDebugFn(fn *ir.Func) bool {

src/cmd/compile/internal/ir/dump.go

@@ -9,11 +9,16 @@
 package ir
 
 import (
+	"crypto/sha256"
+	"encoding/hex"
 	"fmt"
 	"io"
+	"net/url"
 	"os"
 	"reflect"
 	"regexp"
+	"strings"
+	"sync"
 
 	"cmd/compile/internal/base"
 	"cmd/compile/internal/types"
@@ -63,6 +68,127 @@ func FDumpAny(w io.Writer, root any, filter string, depth int) {
 	p.printf("\n")
 }
 
+// MatchAstDump returns true if the fn matches the value
+// of the astdump debug flag.  Fn matches in the following
+// cases:
+//
+//   - astdump == name(fn)
+//   - astdump == pkgname(fn).name(fn)
+//   - astdump == afterslash(pkgname(fn)).name(fn)
+//   - astdump begins with a "~" and what follows "~" is a
+//     regular expression matching pkgname(fn).name(fn)
+//
+// If MatchAstDump returns true, it also prints to os.Stderr
+//
+//	\nir.Match(<fn>, <astdump>) for <where>\n
+func MatchAstDump(fn *Func, where string) bool {
+	if len(base.Debug.AstDump) == 0 {
+		return false
+	}
+	return matchForDump(fn, base.Ctxt.Pkgpath, where)
+}
+
+var dbgRE *regexp.Regexp
+var onceDbgRE sync.Once
+
+func matchForDump(fn *Func, pkgPath, where string) bool {
+	dbg := false
+	flag := base.Debug.AstDump
+	if flag[0] == '~' {
+		onceDbgRE.Do(func() { dbgRE = regexp.MustCompile(flag[1:]) })
+		dbg = dbgRE.MatchString(pkgPath + "." + FuncName(fn))
+	} else {
+		dbg = matchPkgFn(pkgPath, FuncName(fn), flag)
+	}
+	return dbg
+}
+
+// matchPkgFn returns true if pkg and fnName "match" toMatch.
+// "aFunc" matches "aFunc" (in any package)
+// "aPkg.aFunc" matches "aPkg.aFunc"
+// "aPkg/subPkg.aFunc" matches "subPkg.aFunc"
+func matchPkgFn(pkgName, fnName, toMatch string) bool {
+	if fnName == toMatch {
+		return true
+	}
+	matchPkgDotName := func(pkg string) bool {
+		// Allocation-free equality check for toMatch == base.Ctxt.Pkgpath + "." + fnName
+		return len(toMatch) == len(pkg)+1+len(fnName) &&
+			strings.HasPrefix(toMatch, pkg) && toMatch[len(pkg)] == '.' && strings.HasSuffix(toMatch, fnName)
+	}
+	if matchPkgDotName(pkgName) {
+		return true
+	}
+	if l := strings.LastIndexByte(pkgName, '/'); l > 0 && matchPkgDotName(pkgName[l+1:]) {
+		return true
+	}
+
+	return false
+}
+
+// AstDump appends the ast dump for fn to the ast dump file for fn.
+// The generated file name is
+//
+//	url.PathEscape(PkgFuncName(fn)) + ".ast"
+//
+// It also prints
+//
+//	Writing ast output to <astfilename>\n
+//
+// to os.Stderr.
+func AstDump(fn *Func, why string) {
+	err := withLockAndFile(
+		fn,
+		func(w io.Writer) {
+			FDump(w, why, fn)
+		},
+	)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Dump returned error %v\n", err)
+	}
+}
+
+var mu sync.Mutex
+var astDumpFiles = make(map[string]bool)
+
+// escapedFileName constructs a file name from fn and suffix,
+// url-path-escaping the function part of the name and replacing it
+// with a hash if it is too long.  The suffix is neither escaped
+// nor including in the length calculation, so an excessively
+// creative suffix will result in problems.
+func escapedFileName(fn *Func, suffix string) string {
+	name := url.PathEscape(PkgFuncName(fn))
+	if len(name) > 125 { // arbitrary limit on file names, as if anyone types these in by hand
+		hash := sha256.Sum256([]byte(name))
+		name = hex.EncodeToString(hash[:8])
+	}
+	return name + suffix
+}
+
+// withLockAndFile manages ast dump files for various function names
+// and invokes a dumping function to write output, under a lock.
+func withLockAndFile(fn *Func, dump func(io.Writer)) (err error) {
+	name := escapedFileName(fn, ".ast")
+
+	// Ensure that debugging output is not scrambled and is written promptly
+	mu.Lock()
+	defer mu.Unlock()
+	mode := os.O_APPEND | os.O_RDWR
+	if !astDumpFiles[name] {
+		astDumpFiles[name] = true
+		mode = os.O_CREATE | os.O_TRUNC | os.O_RDWR
+		fmt.Fprintf(os.Stderr, "Writing text ast output for %s to %s\n", PkgFuncName(fn), name)
+	}
+
+	fi, err := os.OpenFile(name, mode, 0777)
+	if err != nil {
+		return err
+	}
+	defer func() { err = fi.Close() }()
+	dump(fi)
+	return
+}
+
 type dumper struct {
 	output  io.Writer
 	fieldrx *regexp.Regexp  // field name filter

src/cmd/compile/internal/ir/dump_test.go

@@ -0,0 +1,44 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package ir
+
+import (
+	"testing"
+)
+
+func testMatch(t *testing.T, pkgName, fnName, toMatch string, match bool) {
+	if matchPkgFn(pkgName, fnName, toMatch) != match {
+		t.Errorf("%v != matchPkgFn(%s, %s, %s)", match, pkgName, fnName, toMatch)
+	}
+}
+
+func TestMatchPkgFn(t *testing.T) {
+	// "aFunc" matches "aFunc" (in any package)
+	// "aPkg.aFunc" matches "aPkg.aFunc"
+	// "aPkg/subPkg.aFunc" matches "subPkg.aFunc"
+
+	match := func(pkgName, fnName, toMatch string) {
+		if !matchPkgFn(pkgName, fnName, toMatch) {
+			t.Errorf("matchPkgFn(%s, %s, %s) did not match", pkgName, fnName, toMatch)
+		}
+	}
+	match("aPkg", "AFunc", "AFunc")
+	match("aPkg", "AFunc", "AFunc")
+	match("aPkg", "AFunc", "aPkg.AFunc")
+	match("aPkg/sPkg", "AFunc", "aPkg/sPkg.AFunc")
+	match("aPkg/sPkg", "AFunc", "sPkg.AFunc")
+
+	notmatch := func(pkgName, fnName, toMatch string) {
+		if matchPkgFn(pkgName, fnName, toMatch) {
+			t.Errorf("matchPkgFn(%s, %s, %s) should not match", pkgName, fnName, toMatch)
+		}
+	}
+	notmatch("aPkg", "AFunc", "BFunc")
+	notmatch("aPkg", "AFunc", "aPkg.BFunc")
+	notmatch("aPkg", "AFunc", "bPkg.AFunc")
+	notmatch("aPkg", "AFunc", "aPkg_AFunc")
+	notmatch("aPkg/sPkg", "AFunc", "aPkg/ssPkg.AFunc")
+	notmatch("aPkg/sPkg", "AFunc", "XPkg.AFunc")
+}

src/cmd/compile/internal/ir/fmt.go

@@ -897,11 +897,19 @@ func (l Nodes) Format(s fmt.State, verb rune) {
 // Dump
 
 // Dump prints the message s followed by a debug dump of n.
+// This includes all the recursive structure under n.
 func Dump(s string, n Node) {
 	fmt.Printf("%s%+v\n", s, n)
 }
 
+// Fdump prints to w the message s followed by a debug dump of n.
+// This includes all the recursive structure under n.
+func FDump(w io.Writer, s string, n Node) {
+	fmt.Fprintf(w, "%s%+v\n", s, n)
+}
+
 // DumpList prints the message s followed by a debug dump of each node in the list.
+// This includes all the recursive structure under each node in the list.
 func DumpList(s string, list Nodes) {
 	var buf bytes.Buffer
 	FDumpList(&buf, s, list)
@@ -909,6 +917,7 @@ func DumpList(s string, list Nodes) {
 }
 
 // FDumpList prints to w the message s followed by a debug dump of each node in the list.
+// This includes all the recursive structure under each node in the list.
 func FDumpList(w io.Writer, s string, list Nodes) {
 	io.WriteString(w, s)
 	dumpNodes(w, list, 1)

src/cmd/compile/internal/ir/func.go

@@ -315,7 +315,9 @@ func PkgFuncName(f *Func) string {
 	}
 	s := f.Sym()
 	pkg := s.Pkg
-
+	if pkg == nil {
+		return "<nil>." + s.Name
+	}
 	return pkg.Path + "." + s.Name
 }
 

src/cmd/compile/internal/ir/node.go

@@ -18,6 +18,9 @@ import (
 // A Node is the abstract interface to an IR node.
 type Node interface {
 	// Formatting
+	// For debugging output, use one of
+	//  Dump/FDump/DumpList/FDumplist (in fmt.go)
+	//  DumpAny/FDumpAny (in dump.go)
 	Format(s fmt.State, verb rune)
 
 	// Source position.

src/cmd/compile/internal/loopvar/loopvar.go

@@ -448,6 +448,11 @@ func ForCapture(fn *ir.Func) []VarAndLoop {
 		}
 	}
 	ir.WithFunc(fn, forCapture)
+
+	if ir.MatchAstDump(fn, "loopvar") {
+		ir.AstDump(fn, "loopvar, "+ir.FuncName(fn))
+	}
+
 	return transformed
 }
 

src/cmd/compile/internal/slice/slice.go

@@ -124,6 +124,11 @@ func Funcs(all []*ir.Func) {
 	for _, fn := range all {
 		analyze(fn)
 	}
+	for _, fn := range all {
+		if ir.MatchAstDump(fn, "slice") {
+			ir.AstDump(fn, "slice, "+ir.FuncName(fn))
+		}
+	}
 }
 
 func analyze(fn *ir.Func) {