Make the callgraph construction algorithm configurable

compiler/cmd/coroc/main.go

@@ -19,20 +19,25 @@ USAGE:
   coroc [OPTIONS] [PATH]
 
 OPTIONS:
-  -h, --help      Show this help information
-  -l, --list      List all files that would be compiled
-  -v, --version   Show the compiler version
+  -h, --help         Show this help information
+  -l, --list         List all files that would be compiled
+  -v, --version      Show the compiler version
 
 ADVANCED OPTIONS:
-  -colors         Print debug information about function colors
-  -cpuprofile     Write CPU profile to file
-  -memprofile     Write memory profile to file
+  -callgraph <TYPE>  Set the callgraph construction algorithm
+                     (static, cha, rta, vta). Default is vta.
+
+  -colors            Print debug information about function colors
+
+  -cpuprofile        Write CPU profile to file
+  -memprofile        Write memory profile to file
 `
 
 var (
 	showVersion   bool
 	onlyListFiles bool
 	debugColors   bool
+	callgraphType string
 	cpuProfile    string
 	memProfile    string
 )
@@ -56,6 +61,7 @@ func run() error {
 	boolFlag(&showVersion, "v", "version")
 	boolFlag(&onlyListFiles, "l", "list")
 	boolFlag(&debugColors, "colors")
+	flag.StringVar(&callgraphType, "callgraph", "", "")
 	flag.StringVar(&cpuProfile, "cpuprofile", "", "")
 	flag.StringVar(&memProfile, "memprofile", "", "")
 	flag.Parse()
@@ -109,6 +115,7 @@ func run() error {
 	}
 
 	return compiler.Compile(path,
+		compiler.CallgraphType(callgraphType),
 		compiler.OnlyListFiles(onlyListFiles),
 		compiler.DebugColors(debugColors),
 	)

compiler/compile.go

@@ -17,7 +17,10 @@ import (
 	"strings"
 
 	"golang.org/x/tools/go/ast/astutil"
+	"golang.org/x/tools/go/callgraph"
 	"golang.org/x/tools/go/callgraph/cha"
+	"golang.org/x/tools/go/callgraph/rta"
+	"golang.org/x/tools/go/callgraph/static"
 	"golang.org/x/tools/go/callgraph/vta"
 	"golang.org/x/tools/go/packages"
 	"golang.org/x/tools/go/ssa"
@@ -45,6 +48,7 @@ func Compile(path string, options ...Option) error {
 }
 
 type compiler struct {
+	callgraphType string
 	onlyListFiles bool
 	debugColors   bool
 
@@ -117,10 +121,32 @@ func (c *compiler) compile(path string) error {
 	log.Printf("building SSA program")
 	c.prog, _ = ssautil.AllPackages(pkgs, ssa.InstantiateGenerics|ssa.GlobalDebug)
 	c.prog.Build()
-
-	log.Printf("building call graph")
 	functions := ssautil.AllFunctions(c.prog)
-	cg := vta.CallGraph(functions, cha.CallGraph(c.prog))
+
+	if c.callgraphType == "" {
+		c.callgraphType = "vta"
+	}
+	log.Printf("building callgraph using %s algorithm", c.callgraphType)
+	var cg *callgraph.Graph
+	// See https://cs.opensource.google/go/x/tools/+/refs/tags/v0.16.1:cmd/callgraph/main.go
+	switch c.callgraphType {
+	case "static":
+		cg = static.CallGraph(c.prog)
+	case "cha":
+		cg = cha.CallGraph(c.prog)
+	case "vta":
+		cg = vta.CallGraph(functions, cha.CallGraph(c.prog))
+	case "rta":
+		mains := ssautil.MainPackages(c.prog.AllPackages())
+		var roots []*ssa.Function
+		for _, main := range mains {
+			roots = append(roots, main.Func("init"), main.Func("main"))
+		}
+		rtares := rta.Analyze(roots, true)
+		cg = rtares.CallGraph
+	default:
+		return fmt.Errorf("invalid or unsupported callgraph construction algorithm %q", c.callgraphType)
+	}
 
 	log.Printf("collecting generic instances")
 	c.generics = map[*ssa.Function][]*ssa.Function{}

compiler/options.go

@@ -3,6 +3,12 @@ package compiler
 // Option configures the compiler.
 type Option func(*compiler)
 
+func CallgraphType(callgraphType string) Option {
+	return func(c *compiler) {
+		c.callgraphType = callgraphType
+	}
+}
+
 func OnlyListFiles(enabled bool) Option {
 	return func(c *compiler) {
 		c.onlyListFiles = enabled