cmd/callgraph: a utility for dumping the callgraph of a Go program.

(This functionality is provided by the oracle, but its output
format is inflexible, and the functionality is better suited
to a shell utility.  I may remove the oracle 'callgraph' feature.)

See Usage for details.

+ Test.

LGTM=sameer
R=sameer
CC=golang-codereviews, gri
https://golang.org/cl/164460044

cmd/callgraph/main.go

@@ -0,0 +1,310 @@
+// Copyright 2014 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.
+
+// callgraph: a tool for reporting the call graph of a Go program.
+// See Usage for details, or run with -help.
+package main
+
+// TODO(adonovan):
+//
+// Features:
+// - restrict graph to a single package
+// - output
+//   - functions reachable from root (use digraph tool?)
+//   - unreachable functions (use digraph tool?)
+//   - dynamic (runtime) types
+//   - indexed output (numbered nodes)
+//   - JSON output
+//   - additional template fields:
+//     callee file/line/col
+
+import (
+	"bytes"
+	"flag"
+	"fmt"
+	"go/build"
+	"go/token"
+	"io"
+	"os"
+	"runtime"
+	"text/template"
+
+	"golang.org/x/tools/go/callgraph"
+	"golang.org/x/tools/go/loader"
+	"golang.org/x/tools/go/pointer"
+	"golang.org/x/tools/go/rta"
+	"golang.org/x/tools/go/ssa"
+)
+
+var algoFlag = flag.String("algo", "rta",
+	`Call graph construction algorithm, one of "rta" or "pta"`)
+
+var testFlag = flag.Bool("test", false,
+	"Loads test code (*_test.go) for imported packages")
+
+var formatFlag = flag.String("format",
+	"{{.Caller}}\t--{{.Dynamic}}-{{.Line}}:{{.Column}}-->\t{{.Callee}}",
+	"A template expression specifying how to format an edge")
+
+const Usage = `callgraph: display the the call graph of a Go program.
+
+Usage:
+
+  callgraph [-algo=rta|pta] [-test] [-format=...] <args>...
+
+Flags:
+
+-algo      Specifies the call-graph construction algorithm.  One of:
+           "rta": Rapid Type Analysis  (simple and fast)
+           "pta": inclusion-based Points-To Analysis (slower but more precise)
+
+-test      Include the package's tests in the analysis.
+
+-format    Specifies the format in which each call graph edge is displayed.
+           One of:
+           "digraph":   output suitable for input to
+                        golang.org/x/tools/cmd/digraph.
+           "graphviz":  output in AT&T GraphViz (.dot) format.
+
+           All other values are interpreted using text/template syntax.
+           The default value is:
+
+            {{.Caller}}\t--{{.Dynamic}}-{{.Line}}:{{.Column}}-->\t{{.Callee}}
+
+           The structure passed to the template is (effectively):
+
+                   type Edge struct {
+                           Caller      *ssa.Function // calling function
+                           Callee      *ssa.Function // called function
+
+                           // Call site:
+                           Filename    string // containing file
+                           Offset      int    // offset within file of '('
+                           Line        int    // line number
+                           Column      int    // column number of call
+                           Dynamic     string // "static" or "dynamic"
+                           Description string // e.g. "static method call"
+                   }
+
+           Caller and Callee are *ssa.Function values, which print as
+           "(*sync/atomic.Mutex).Lock", but other attributes may be
+           derived from them, e.g. Caller.Pkg.Object.Path yields the
+           import path of the enclosing package.  Consult the go/ssa
+           API documentation for details.
+
+` + loader.FromArgsUsage + `
+
+Examples:
+
+  Show the call graph of the trivial web server application:
+
+    callgraph -format digraph $GOROOT/src/net/http/triv.go
+
+  Same, but show only the packages of each function:
+
+    callgraph -format '{{.Caller.Pkg.Object.Path}} -> {{.Callee.Pkg.Object.Path}}' \
+      $GOROOT/src/net/http/triv.go | sort | uniq
+
+  Show functions that make dynamic calls into the 'fmt' test package,
+  using the pointer analysis algorithm:
+
+    callgraph -format='{{.Caller}} -{{.Dynamic}}-> {{.Callee}}' -test -algo=pta fmt |
+      sed -ne 's/-dynamic-/--/p' |
+      sed -ne 's/-->.*fmt_test.*$//p' | sort | uniq
+
+  Show all functions directly called by the callgraph tool's main function:
+
+    callgraph -format=digraph golang.org/x/tools/cmd/callgraph |
+      digraph succs golang.org/x/tools/cmd/callgraph.main
+`
+
+func init() {
+	// If $GOMAXPROCS isn't set, use the full capacity of the machine.
+	// For small machines, use at least 4 threads.
+	if os.Getenv("GOMAXPROCS") == "" {
+		n := runtime.NumCPU()
+		if n < 4 {
+			n = 4
+		}
+		runtime.GOMAXPROCS(n)
+	}
+}
+
+func main() {
+	flag.Parse()
+	if err := doCallgraph(&build.Default, *algoFlag, *formatFlag, *testFlag, flag.Args()); err != nil {
+		fmt.Fprintf(os.Stderr, "callgraph: %s.\n", err)
+		os.Exit(1)
+	}
+}
+
+var stdout io.Writer = os.Stdout
+
+func doCallgraph(ctxt *build.Context, algo, format string, tests bool, args []string) error {
+	conf := loader.Config{
+		Build:         ctxt,
+		SourceImports: true,
+	}
+
+	if len(args) == 0 {
+		fmt.Fprintln(os.Stderr, Usage)
+		return nil
+	}
+
+	// Use the initial packages from the command line.
+	args, err := conf.FromArgs(args, tests)
+	if err != nil {
+		return err
+	}
+
+	// Load, parse and type-check the whole program.
+	iprog, err := conf.Load()
+	if err != nil {
+		return err
+	}
+
+	// Create and build SSA-form program representation.
+	prog := ssa.Create(iprog, 0)
+	prog.BuildAll()
+
+	// Determine the main package.
+	// TODO(adonovan): allow independent control over tests, mains
+	// and libraries.
+	// TODO(adonovan): put this logic in a library; we keep reinventing it.
+	var main *ssa.Package
+	pkgs := prog.AllPackages()
+	if tests {
+		// If -test, use all packages' tests.
+		if len(pkgs) > 0 {
+			main = prog.CreateTestMainPackage(pkgs...)
+		}
+		if main == nil {
+			return fmt.Errorf("no tests")
+		}
+	} else {
+		// Otherwise, use main.main.
+		for _, pkg := range pkgs {
+			if pkg.Object.Name() == "main" {
+				main = pkg
+				if main.Func("main") == nil {
+					return fmt.Errorf("no func main() in main package")
+				}
+				break
+			}
+		}
+		if main == nil {
+			return fmt.Errorf("no main package")
+		}
+	}
+
+	// Invariant: main package has a main() function.
+
+	// -- call graph construction ------------------------------------------
+
+	var cg *callgraph.Graph
+
+	switch algo {
+	case "pta":
+		config := &pointer.Config{
+			Mains:          []*ssa.Package{main},
+			BuildCallGraph: true,
+		}
+		ptares, err := pointer.Analyze(config)
+		if err != nil {
+			return err // internal error in pointer analysis
+		}
+		cg = ptares.CallGraph
+
+	case "rta":
+		roots := []*ssa.Function{
+			main.Func("init"),
+			main.Func("main"),
+		}
+		rtares := rta.Analyze(roots, true)
+		cg = rtares.CallGraph
+
+		// NB: RTA gives us Reachable and RuntimeTypes too.
+
+	default:
+		return fmt.Errorf("unknown algorithm: %s", algo)
+	}
+
+	cg.DeleteSyntheticNodes()
+
+	// -- output------------------------------------------------------------
+
+	var before, after string
+
+	// Pre-canned formats.
+	switch format {
+	case "digraph":
+		format = `{{printf "%q %q" .Caller .Callee}}`
+
+	case "graphviz":
+		before = "digraph callgraph {\n"
+		after = "}\n"
+		format = `  {{printf "%q" .Caller}} -> {{printf "%q" .Callee}}"`
+	}
+
+	tmpl, err := template.New("-format").Parse(format)
+	if err != nil {
+		return fmt.Errorf("invalid -format template: %v", err)
+	}
+
+	// Allocate these once, outside the traversal.
+	var buf bytes.Buffer
+	data := Edge{fset: prog.Fset}
+
+	fmt.Fprint(stdout, before)
+	if err := callgraph.GraphVisitEdges(cg, func(edge *callgraph.Edge) error {
+		data.position.Offset = -1
+		data.edge = edge
+		data.Caller = edge.Caller.Func
+		data.Callee = edge.Callee.Func
+
+		buf.Reset()
+		if err := tmpl.Execute(&buf, &data); err != nil {
+			return err
+		}
+		stdout.Write(buf.Bytes())
+		if len := buf.Len(); len == 0 || buf.Bytes()[len-1] != '\n' {
+			fmt.Fprintln(stdout)
+		}
+		return nil
+	}); err != nil {
+		return err
+	}
+	fmt.Fprint(stdout, after)
+	return nil
+}
+
+type Edge struct {
+	Caller *ssa.Function
+	Callee *ssa.Function
+
+	edge     *callgraph.Edge
+	fset     *token.FileSet
+	position token.Position // initialized lazily
+}
+
+func (e *Edge) pos() *token.Position {
+	if e.position.Offset == -1 {
+		e.position = e.fset.Position(e.edge.Pos()) // called lazily
+	}
+	return &e.position
+}
+
+func (e *Edge) Filename() string { return e.pos().Filename }
+func (e *Edge) Column() int      { return e.pos().Column }
+func (e *Edge) Line() int        { return e.pos().Line }
+func (e *Edge) Offset() int      { return e.pos().Offset }
+
+func (e *Edge) Dynamic() string {
+	if e.edge.Site != nil && e.edge.Site.Common().StaticCallee() == nil {
+		return "dynamic"
+	}
+	return "static"
+}
+
+func (e *Edge) Description() string { return e.edge.Description() }

cmd/callgraph/main_test.go

@@ -0,0 +1,73 @@
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"go/build"
+	"reflect"
+	"sort"
+	"strings"
+	"testing"
+)
+
+func TestCallgraph(t *testing.T) {
+	ctxt := build.Default // copy
+	ctxt.GOPATH = "testdata"
+
+	const format = "{{.Caller}} --> {{.Callee}}"
+
+	for _, test := range []struct {
+		algo, format string
+		tests        bool
+		want         []string
+	}{
+		{"rta", format, false, []string{
+			// rta imprecisely shows cross product of {main,main2} x {C,D}
+			`pkg.main --> (pkg.C).f`,
+			`pkg.main --> (pkg.D).f`,
+			`pkg.main --> pkg.main2`,
+			`pkg.main2 --> (pkg.C).f`,
+			`pkg.main2 --> (pkg.D).f`,
+		}},
+		{"pta", format, false, []string{
+			// pta distinguishes main->C, main2->D.  Also has a root node.
+			`<root> --> pkg.init`,
+			`<root> --> pkg.main`,
+			`pkg.main --> (pkg.C).f`,
+			`pkg.main --> pkg.main2`,
+			`pkg.main2 --> (pkg.D).f`,
+		}},
+		// tests: main is not called.
+		{"rta", format, true, []string{
+			`pkg.Example --> (pkg.C).f`,
+			`test$main.init --> pkg.init`,
+		}},
+		{"pta", format, true, []string{
+			`<root> --> pkg.Example`,
+			`<root> --> test$main.init`,
+			`pkg.Example --> (pkg.C).f`,
+			`test$main.init --> pkg.init`,
+		}},
+	} {
+		stdout = new(bytes.Buffer)
+		if err := doCallgraph(&ctxt, test.algo, test.format, test.tests, []string{"pkg"}); err != nil {
+			t.Error(err)
+			continue
+		}
+
+		got := sortedLines(fmt.Sprint(stdout))
+		if !reflect.DeepEqual(got, test.want) {
+			t.Errorf("callgraph(%q, %q, %t):\ngot:\n%s\nwant:\n%s",
+				test.algo, test.format, test.tests,
+				strings.Join(got, "\n"),
+				strings.Join(test.want, "\n"))
+		}
+	}
+}
+
+func sortedLines(s string) []string {
+	s = strings.TrimSpace(s)
+	lines := strings.Split(s, "\n")
+	sort.Strings(lines)
+	return lines
+}

cmd/callgraph/testdata/src/pkg/pkg.go

@@ -0,0 +1,25 @@
+package main
+
+type I interface {
+	f()
+}
+
+type C int
+
+func (C) f() {}
+
+type D int
+
+func (D) f() {}
+
+func main() {
+	var i I = C(0)
+	i.f() // dynamic call
+
+	main2()
+}
+
+func main2() {
+	var i I = D(0)
+	i.f() // dynamic call
+}

cmd/callgraph/testdata/src/pkg/pkg_test.go

@@ -0,0 +1,7 @@
+package main
+
+// Don't import "testing", it adds a lot of callgraph edges.
+
+func Example() {
+	C(0).f()
+}