go/types/typeutil: Callee supports generics

Adds support for finding the Callee if it is a generic function.
Generic methods happen to already work without this change.
Updating this so that they are consistent.

For golang/go#48704

Change-Id: I649ec746e350db4a0086ed31535b2e14baa32314
Reviewed-on: https://go-review.googlesource.com/c/tools/+/359974
Run-TryBot: Tim King <taking@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Zvonimir Pavlinovic <zpavlinovic@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>
Trust: Zvonimir Pavlinovic <zpavlinovic@google.com>

Author: timothy-king

go/types/typeutil/callee.go

@@ -9,13 +9,30 @@ import (
 	"go/types"
 
 	"golang.org/x/tools/go/ast/astutil"
+	"golang.org/x/tools/internal/typeparams"
 )
 
 // Callee returns the named target of a function call, if any:
 // a function, method, builtin, or variable.
+//
+// Functions and methods may potentially have type parameters.
 func Callee(info *types.Info, call *ast.CallExpr) types.Object {
+	fun := astutil.Unparen(call.Fun)
+
+	// Look through type instantiation if necessary.
+	isInstance := false
+	switch fun.(type) {
+	case *ast.IndexExpr, *typeparams.IndexListExpr:
+		// When extracting the callee from an *IndexExpr, we need to check that
+		// it is a *types.Func and not a *types.Var.
+		// Example: Don't match a slice m within the expression `m[0]()`.
+		isInstance = true
+		ix := typeparams.GetIndexExprData(fun)
+		fun = ix.X
+	}
+
 	var obj types.Object
-	switch fun := astutil.Unparen(call.Fun).(type) {
+	switch fun := fun.(type) {
 	case *ast.Ident:
 		obj = info.Uses[fun] // type, var, builtin, or declared func
 	case *ast.SelectorExpr:
@@ -28,11 +45,18 @@ func Callee(info *types.Info, call *ast.CallExpr) types.Object {
 	if _, ok := obj.(*types.TypeName); ok {
 		return nil // T(x) is a conversion, not a call
 	}
+	// A Func is required to match instantiations.
+	if _, ok := obj.(*types.Func); isInstance && !ok {
+		return nil // Was not a Func.
+	}
 	return obj
 }
 
-// StaticCallee returns the target (function or method) of a static
-// function call, if any. It returns nil for calls to builtins.
+// StaticCallee returns the target (function or method) of a static function
+// call, if any. It returns nil for calls to builtins.
+//
+// Note: for calls of instantiated functions and methods, StaticCallee returns
+// the corresponding generic function or method on the generic type.
 func StaticCallee(info *types.Info, call *ast.CallExpr) *types.Func {
 	if f, ok := Callee(info, call).(*types.Func); ok && !interfaceMethod(f) {
 		return f

go/types/typeutil/callee_test.go

@@ -5,85 +5,167 @@
 package typeutil_test
 
 import (
+	"fmt"
 	"go/ast"
-	"go/importer"
 	"go/parser"
 	"go/token"
 	"go/types"
 	"strings"
 	"testing"
 
 	"golang.org/x/tools/go/types/typeutil"
+	"golang.org/x/tools/internal/typeparams"
 )
 
 func TestStaticCallee(t *testing.T) {
-	const src = `package p
+	testStaticCallee(t, []string{
+		`package q;
+		func Abs(x int) int {
+			if x < 0 {
+				return -x
+			}
+			return x
+		}`,
+		`package p
+		import "q"
 
-import "fmt"
+		type T int
 
-type T int
+		func g(int)
 
-func g(int)
+		var f = g
 
-var f = g
+		var x int
 
-var x int
+		type s struct{ f func(int) }
+		func (s) g(int)
 
-type s struct{ f func(int) }
-func (s) g(int)
+		type I interface{ f(int) }
 
-type I interface{ f(int) }
+		var a struct{b struct{c s}}
 
-var a struct{b struct{c s}}
+		var n map[int]func()
+		var m []func()
 
-func calls() {
-	g(x)           // a declared func
-	s{}.g(x)       // a concrete method
-	a.b.c.g(x)     // same
-	fmt.Println(x) // declared func, qualified identifier
-}
+		func calls() {
+			g(x)           // a declared func
+			s{}.g(x)       // a concrete method
+			a.b.c.g(x)     // same
+			_ = q.Abs(x)   // declared func, qualified identifier
+		}
 
-func noncalls() {
-	_ = T(x)    // a type
-	f(x)        // a var
-	panic(x)    // a built-in
-	s{}.f(x)    // a field
-	I(nil).f(x) // interface method
+		func noncalls() {
+			_ = T(x)    // a type
+			f(x)        // a var
+			panic(x)    // a built-in
+			s{}.f(x)    // a field
+			I(nil).f(x) // interface method
+			m[0]()      // a map
+			n[0]()      // a slice
+		}
+		`})
 }
-`
-	// parse
-	fset := token.NewFileSet()
-	f, err := parser.ParseFile(fset, "p.go", src, 0)
-	if err != nil {
-		t.Fatal(err)
+
+func TestTypeParamStaticCallee(t *testing.T) {
+	if !typeparams.Enabled {
+		t.Skip("type parameters are not enabled")
 	}
+	testStaticCallee(t, []string{
+		`package q
+		func R[T any]() {}
+		`,
+		`package p
+		import "q"
+		type I interface{
+			i()
+		}
+
+		type G[T any] func() T
+		func F[T any]() T { var x T; return x }
+
+		type M[T I] struct{ t T }
+		func (m M[T]) noncalls() {
+			m.t.i()   // method on a type parameter
+		}
+
+		func (m M[T]) calls() {
+			m.calls() // method on a generic type
+		}
+
+		type Chain[T I] struct{ r struct { s M[T] } }
 
-	// type-check
+		type S int
+		func (S) i() {}
+
+		func Multi[TP0, TP1 any](){}
+
+		func calls() {
+			_ = F[int]()            // instantiated function
+			_ = (F[int])()          // go through parens
+			M[S]{}.calls()          // instantiated method
+			Chain[S]{}.r.s.calls()  // same as above
+			Multi[int,string]()     // multiple type parameters
+			q.R[int]()              // different package
+		}
+
+		func noncalls() {
+			_ = G[int](nil)()  // instantiated function
+		}
+		`})
+}
+
+// testStaticCallee parses and type checks each file content in contents
+// as a single file package in order. Within functions that have the suffix
+// "calls" it checks that the CallExprs within have a static callee.
+// If the function's name == "calls" all calls must have static callees,
+// and if the name != "calls", the calls must not have static callees.
+// Failures are reported on t.
+func testStaticCallee(t *testing.T, contents []string) {
+	fset := token.NewFileSet()
+	packages := make(map[string]*types.Package)
+	cfg := &types.Config{Importer: closure(packages)}
 	info := &types.Info{
 		Uses:       make(map[*ast.Ident]types.Object),
 		Selections: make(map[*ast.SelectorExpr]*types.Selection),
 	}
-	cfg := &types.Config{Importer: importer.ForCompiler(fset, "source", nil)}
-	if _, err := cfg.Check("p", fset, []*ast.File{f}, info); err != nil {
-		t.Fatal(err)
+	typeparams.InitInstanceInfo(info)
+
+	var files []*ast.File
+	for i, content := range contents {
+		// parse
+		f, err := parser.ParseFile(fset, fmt.Sprintf("%d.go", i), content, 0)
+		if err != nil {
+			t.Fatal(err)
+		}
+		files = append(files, f)
+
+		// type-check
+		pkg, err := cfg.Check(f.Name.Name, fset, []*ast.File{f}, info)
+		if err != nil {
+			t.Fatal(err)
+		}
+		packages[pkg.Path()] = pkg
 	}
 
-	for _, decl := range f.Decls {
-		if decl, ok := decl.(*ast.FuncDecl); ok && strings.HasSuffix(decl.Name.Name, "calls") {
-			wantCallee := decl.Name.Name == "calls" // false within func noncalls()
-			ast.Inspect(decl.Body, func(n ast.Node) bool {
-				if call, ok := n.(*ast.CallExpr); ok {
-					fn := typeutil.StaticCallee(info, call)
-					if fn == nil && wantCallee {
-						t.Errorf("%s: StaticCallee returned nil",
-							fset.Position(call.Lparen))
-					} else if fn != nil && !wantCallee {
-						t.Errorf("%s: StaticCallee returned %s, want nil",
-							fset.Position(call.Lparen), fn)
+	// check
+	for _, f := range files {
+		for _, decl := range f.Decls {
+			if decl, ok := decl.(*ast.FuncDecl); ok && strings.HasSuffix(decl.Name.Name, "calls") {
+				wantCallee := decl.Name.Name == "calls" // false within func noncalls()
+				ast.Inspect(decl.Body, func(n ast.Node) bool {
+					if call, ok := n.(*ast.CallExpr); ok {
+						fn := typeutil.StaticCallee(info, call)
+						if fn == nil && wantCallee {
+							t.Errorf("%s: StaticCallee returned nil",
+								fset.Position(call.Lparen))
+						} else if fn != nil && !wantCallee {
+							t.Errorf("%s: StaticCallee returned %s, want nil",
+								fset.Position(call.Lparen), fn)
+						}
 					}
-				}
-				return true
-			})
+					return true
+				})
+			}
 		}
 	}
 }