go/types, types2: implement Info.PkgNameOf

For #62037.

Change-Id: I354f6417232708278d3f2b2d5ea41ff48e08d6b6
Reviewed-on: https://go-review.googlesource.com/c/go/+/541575
Reviewed-by: Alan Donovan <adonovan@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Robert Griesemer <gri@google.com>
Auto-Submit: Robert Griesemer <gri@google.com>
Run-TryBot: Robert Griesemer <gri@google.com>

5 files changed, 185 insertions, 0 deletions

diff --git a/api/next/62037.txt b/api/next/62037.txt
new file mode 100644
index 0000000000..78374214c8
--- /dev/null
+++ b/api/next/62037.txt
@@ -0,0 +1,2 @@
+pkg go/types, method (*Info) PkgNameOf(*ast.ImportSpec) *PkgName #62037
+pkg go/types, method (Checker) PkgNameOf(*ast.ImportSpec) *PkgName #62037
diff --git a/src/cmd/compile/internal/types2/api.go b/src/cmd/compile/internal/types2/api.go
index 675882c49f..2093ceb817 100644
--- a/src/cmd/compile/internal/types2/api.go
+++ b/src/cmd/compile/internal/types2/api.go
@@ -342,6 +342,23 @@ func (info *Info) ObjectOf(id *syntax.Name) Object {
 	return info.Uses[id]
 }
 
+// PkgNameOf returns the local package name defined by the import,
+// or nil if not found.
+//
+// For dot-imports, the package name is ".".
+//
+// Precondition: the Defs and Implicts maps are populated.
+func (info *Info) PkgNameOf(imp *syntax.ImportDecl) *PkgName {
+	var obj Object
+	if imp.LocalPkgName != nil {
+		obj = info.Defs[imp.LocalPkgName]
+	} else {
+		obj = info.Implicits[imp]
+	}
+	pkgname, _ := obj.(*PkgName)
+	return pkgname
+}
+
 // TypeAndValue reports the type and value (for constants)
 // of the corresponding expression.
 type TypeAndValue struct {
diff --git a/src/cmd/compile/internal/types2/api_test.go b/src/cmd/compile/internal/types2/api_test.go
index f5bdcf213d..a2621854bc 100644
--- a/src/cmd/compile/internal/types2/api_test.go
+++ b/src/cmd/compile/internal/types2/api_test.go
@@ -958,6 +958,80 @@ func TestImplicitsInfo(t *testing.T) {
 	}
 }
 
+func TestPkgNameOf(t *testing.T) {
+	testenv.MustHaveGoBuild(t)
+
+	const src = `
+package p
+
+import (
+	. "os"
+	_ "io"
+	"math"
+	"path/filepath"
+	snort "sort"
+)
+
+// avoid imported and not used errors
+var (
+	_ = Open // os.Open
+	_ = math.Sin
+	_ = filepath.Abs
+	_ = snort.Ints
+)
+`
+
+	var tests = []struct {
+		path string // path string enclosed in "'s
+		want string
+	}{
+		{`"os"`, "."},
+		{`"io"`, "_"},
+		{`"math"`, "math"},
+		{`"path/filepath"`, "filepath"},
+		{`"sort"`, "snort"},
+	}
+
+	f := mustParse(src)
+	info := Info{
+		Defs:      make(map[*syntax.Name]Object),
+		Implicits: make(map[syntax.Node]Object),
+	}
+	var conf Config
+	conf.Importer = defaultImporter()
+	_, err := conf.Check("p", []*syntax.File{f}, &info)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// map import paths to importDecl
+	imports := make(map[string]*syntax.ImportDecl)
+	for _, d := range f.DeclList {
+		if imp, _ := d.(*syntax.ImportDecl); imp != nil {
+			imports[imp.Path.Value] = imp
+		}
+	}
+
+	for _, test := range tests {
+		imp := imports[test.path]
+		if imp == nil {
+			t.Fatalf("invalid test case: import path %s not found", test.path)
+		}
+		got := info.PkgNameOf(imp)
+		if got == nil {
+			t.Fatalf("import %s: package name not found", test.path)
+		}
+		if got.Name() != test.want {
+			t.Errorf("import %s: got %s; want %s", test.path, got.Name(), test.want)
+		}
+	}
+
+	// test non-existing importDecl
+	if got := info.PkgNameOf(new(syntax.ImportDecl)); got != nil {
+		t.Errorf("got %s for non-existing import declaration", got.Name())
+	}
+}
+
 func predString(tv TypeAndValue) string {
 	var buf strings.Builder
 	pred := func(b bool, s string) {
diff --git a/src/go/types/api.go b/src/go/types/api.go
index 38cde80085..0dc06af538 100644
--- a/src/go/types/api.go
+++ b/src/go/types/api.go
@@ -331,6 +331,23 @@ func (info *Info) ObjectOf(id *ast.Ident) Object {
 	return info.Uses[id]
 }
 
+// PkgNameOf returns the local package name defined by the import,
+// or nil if not found.
+//
+// For dot-imports, the package name is ".".
+//
+// Precondition: the Defs and Implicts maps are populated.
+func (info *Info) PkgNameOf(imp *ast.ImportSpec) *PkgName {
+	var obj Object
+	if imp.Name != nil {
+		obj = info.Defs[imp.Name]
+	} else {
+		obj = info.Implicits[imp]
+	}
+	pkgname, _ := obj.(*PkgName)
+	return pkgname
+}
+
 // TypeAndValue reports the type and value (for constants)
 // of the corresponding expression.
 type TypeAndValue struct {
diff --git a/src/go/types/api_test.go b/src/go/types/api_test.go
index 0769890101..3050b930b5 100644
--- a/src/go/types/api_test.go
+++ b/src/go/types/api_test.go
@@ -960,6 +960,81 @@ func TestImplicitsInfo(t *testing.T) {
 	}
 }
 
+func TestPkgNameOf(t *testing.T) {
+	testenv.MustHaveGoBuild(t)
+
+	const src = `
+package p
+
+import (
+	. "os"
+	_ "io"
+	"math"
+	"path/filepath"
+	snort "sort"
+)
+
+// avoid imported and not used errors
+var (
+	_ = Open // os.Open
+	_ = math.Sin
+	_ = filepath.Abs
+	_ = snort.Ints
+)
+`
+
+	var tests = []struct {
+		path string // path string enclosed in "'s
+		want string
+	}{
+		{`"os"`, "."},
+		{`"io"`, "_"},
+		{`"math"`, "math"},
+		{`"path/filepath"`, "filepath"},
+		{`"sort"`, "snort"},
+	}
+
+	fset := token.NewFileSet()
+	f := mustParse(fset, src)
+	info := Info{
+		Defs:      make(map[*ast.Ident]Object),
+		Implicits: make(map[ast.Node]Object),
+	}
+	var conf Config
+	conf.Importer = importer.Default()
+	_, err := conf.Check("p", fset, []*ast.File{f}, &info)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// map import paths to importDecl
+	imports := make(map[string]*ast.ImportSpec)
+	for _, s := range f.Decls[0].(*ast.GenDecl).Specs {
+		if imp, _ := s.(*ast.ImportSpec); imp != nil {
+			imports[imp.Path.Value] = imp
+		}
+	}
+
+	for _, test := range tests {
+		imp := imports[test.path]
+		if imp == nil {
+			t.Fatalf("invalid test case: import path %s not found", test.path)
+		}
+		got := info.PkgNameOf(imp)
+		if got == nil {
+			t.Fatalf("import %s: package name not found", test.path)
+		}
+		if got.Name() != test.want {
+			t.Errorf("import %s: got %s; want %s", test.path, got.Name(), test.want)
+		}
+	}
+
+	// test non-existing importDecl
+	if got := info.PkgNameOf(new(ast.ImportSpec)); got != nil {
+		t.Errorf("got %s for non-existing import declaration", got.Name())
+	}
+}
+
 func predString(tv TypeAndValue) string {
 	var buf strings.Builder
 	pred := func(b bool, s string) {