1 files changed, 62 insertions, 27 deletions

diff --git a/go/callgraph/cha/cha_test.go b/go/callgraph/cha/cha_test.go
index 3dc03143b..a12b3d0a3 100644
--- a/go/callgraph/cha/cha_test.go
+++ b/go/callgraph/cha/cha_test.go
@@ -24,7 +24,9 @@ import (
 	"golang.org/x/tools/go/callgraph"
 	"golang.org/x/tools/go/callgraph/cha"
 	"golang.org/x/tools/go/loader"
+	"golang.org/x/tools/go/ssa"
 	"golang.org/x/tools/go/ssa/ssautil"
+	"golang.org/x/tools/internal/typeparams"
 )
 
 var inputs = []string{
@@ -47,19 +49,9 @@ func expectation(f *ast.File) (string, token.Pos) {
 // TestCHA runs CHA on each file in inputs, prints the dynamic edges of
 // the call graph, and compares it with the golden results embedded in
 // the WANT comment at the end of the file.
-//
 func TestCHA(t *testing.T) {
 	for _, filename := range inputs {
-		content, err := ioutil.ReadFile(filename)
-		if err != nil {
-			t.Errorf("couldn't read file '%s': %s", filename, err)
-			continue
-		}
-
-		conf := loader.Config{
-			ParserMode: parser.ParseComments,
-		}
-		f, err := conf.ParseFile(filename, content)
+		prog, f, mainPkg, err := loadProgInfo(filename, ssa.InstantiateGenerics)
 		if err != nil {
 			t.Error(err)
 			continue
@@ -67,34 +59,77 @@ func TestCHA(t *testing.T) {
 
 		want, pos := expectation(f)
 		if pos == token.NoPos {
-			t.Errorf("No WANT: comment in %s", filename)
-			continue
-		}
-
-		conf.CreateFromFiles("main", f)
-		iprog, err := conf.Load()
-		if err != nil {
-			t.Error(err)
+			t.Error(fmt.Errorf("No WANT: comment in %s", filename))
 			continue
 		}
 
-		prog := ssautil.CreateProgram(iprog, 0)
-		mainPkg := prog.Package(iprog.Created[0].Pkg)
-		prog.Build()
-
 		cg := cha.CallGraph(prog)
 
-		if got := printGraph(cg, mainPkg.Pkg); got != want {
+		if got := printGraph(cg, mainPkg.Pkg, "dynamic", "Dynamic calls"); got != want {
 			t.Errorf("%s: got:\n%s\nwant:\n%s",
 				prog.Fset.Position(pos), got, want)
 		}
 	}
 }
 
-func printGraph(cg *callgraph.Graph, from *types.Package) string {
+// TestCHAGenerics is TestCHA tailored for testing generics,
+func TestCHAGenerics(t *testing.T) {
+	if !typeparams.Enabled {
+		t.Skip("TestCHAGenerics requires type parameters")
+	}
+
+	filename := "testdata/generics.go"
+	prog, f, mainPkg, err := loadProgInfo(filename, ssa.InstantiateGenerics)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	want, pos := expectation(f)
+	if pos == token.NoPos {
+		t.Fatal(fmt.Errorf("No WANT: comment in %s", filename))
+	}
+
+	cg := cha.CallGraph(prog)
+
+	if got := printGraph(cg, mainPkg.Pkg, "", "All calls"); got != want {
+		t.Errorf("%s: got:\n%s\nwant:\n%s",
+			prog.Fset.Position(pos), got, want)
+	}
+}
+
+func loadProgInfo(filename string, mode ssa.BuilderMode) (*ssa.Program, *ast.File, *ssa.Package, error) {
+	content, err := ioutil.ReadFile(filename)
+	if err != nil {
+		return nil, nil, nil, fmt.Errorf("couldn't read file '%s': %s", filename, err)
+	}
+
+	conf := loader.Config{
+		ParserMode: parser.ParseComments,
+	}
+	f, err := conf.ParseFile(filename, content)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	conf.CreateFromFiles("main", f)
+	iprog, err := conf.Load()
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	prog := ssautil.CreateProgram(iprog, mode)
+	prog.Build()
+
+	return prog, f, prog.Package(iprog.Created[0].Pkg), nil
+}
+
+// printGraph returns a string representation of cg involving only edges
+// whose description contains edgeMatch. The string representation is
+// prefixed with a desc line.
+func printGraph(cg *callgraph.Graph, from *types.Package, edgeMatch string, desc string) string {
 	var edges []string
 	callgraph.GraphVisitEdges(cg, func(e *callgraph.Edge) error {
-		if strings.Contains(e.Description(), "dynamic") {
+		if strings.Contains(e.Description(), edgeMatch) {
 			edges = append(edges, fmt.Sprintf("%s --> %s",
 				e.Caller.Func.RelString(from),
 				e.Callee.Func.RelString(from)))
@@ -104,7 +139,7 @@ func printGraph(cg *callgraph.Graph, from *types.Package) string {
 	sort.Strings(edges)
 
 	var buf bytes.Buffer
-	buf.WriteString("Dynamic calls\n")
+	buf.WriteString(desc + "\n")
 	for _, edge := range edges {
 		fmt.Fprintf(&buf, "  %s\n", edge)
 	}