imports: use go/packages, support modules

When in module mode (go env GOMOD is nonempty), use go/packages to load
package information. This adds module support, without slowing down
non-module mode. Functionality should be the same in both cases.

Since we intend to use go/packages for everything eventually, the tests
run in three modes: GOPATH, using the old code paths; GOPATH_GoPackages,
using go/packages on a GOPATH layout; and Modules, using go/packages on
a module layout. They do take a little while to run.

The code path for go/packages is a lot more convoluted than it should be,
since it mixes in and out of the older path. I expect we'll clean that
up after deleting the old one.

Change-Id: I5b54bd470e091e6fb429597dfbfd1c8fa03f5b59
Reviewed-on: https://go-review.googlesource.com/c/152317
Run-TryBot: Heschi Kreinick <heschi@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Reviewed-by: Ian Cottrell <iancottrell@google.com>

imports/fix.go

@@ -5,6 +5,7 @@
 package imports
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"go/ast"
@@ -14,6 +15,7 @@ import (
 	"io/ioutil"
 	"log"
 	"os"
+	"os/exec"
 	"path"
 	"path/filepath"
 	"sort"
@@ -22,6 +24,7 @@ import (
 	"sync"
 
 	"golang.org/x/tools/go/ast/astutil"
+	"golang.org/x/tools/go/packages"
 	"golang.org/x/tools/internal/gopathwalk"
 )
 
@@ -89,6 +92,8 @@ type packageInfo struct {
 // parseOtherFiles parses all the Go files in srcDir except filename, including
 // test files if filename looks like a test.
 func parseOtherFiles(fset *token.FileSet, srcDir, filename string) []*ast.File {
+	// This could use go/packages but it doesn't buy much, and it fails
+	// with https://golang.org/issue/26296 in LoadFiles mode in some cases.
 	considerTests := strings.HasSuffix(filename, "_test.go")
 
 	fileBase := filepath.Base(filename)
@@ -228,11 +233,12 @@ func (p *pass) findMissingImport(pkg string, syms map[string]bool) *importInfo {
 // It can be modified in some ways during use; see comments below.
 type pass struct {
 	// Inputs. These must be set before a call to load, and not modified after.
-	fset       *token.FileSet                          // fset used to parse f and its siblings.
-	f          *ast.File                               // the file being fixed.
-	otherFiles []*ast.File                             // sibling files.
-	srcDir     string                                  // the directory containing f.
-	pathToName func(path string, srcDir string) string // the function to use to resolve an import path to a package name.
+	fset                 *token.FileSet // fset used to parse f and its siblings.
+	f                    *ast.File      // the file being fixed.
+	srcDir               string         // the directory containing f.
+	useGoPackages        bool           // use go/packages to load package information.
+	loadRealPackageNames bool           // if true, load package names from disk rather than guessing them.
+	otherFiles           []*ast.File    // sibling files.
 
 	// Intermediate state, generated by load.
 	existingImports map[string]*importInfo
@@ -240,23 +246,57 @@ type pass struct {
 	used            map[*importInfo]bool
 
 	// Inputs to fix. These can be augmented between successive fix calls.
-	lastTry        bool                    // indicates that this is the last call and fix should clean up as best it can.
-	addImportNames bool                    // add names to mismatched imports.
-	candidates     []*importInfo           // candidate imports in priority order.
-	knownPackages  map[string]*packageInfo // information about all known packages.
+	lastTry       bool                    // indicates that this is the last call and fix should clean up as best it can.
+	candidates    []*importInfo           // candidate imports in priority order.
+	knownPackages map[string]*packageInfo // information about all known packages.
 }
 
 // loadPackageNames saves the package names for everything referenced by imports.
-func (p *pass) loadPackageNames(imports []*importInfo) {
+func (p *pass) loadPackageNames(imports []*importInfo) error {
+	var unknown []string
 	for _, imp := range imports {
 		if _, ok := p.knownPackages[imp.importPath]; ok {
 			continue
 		}
-		p.knownPackages[imp.importPath] = &packageInfo{
-			name:    p.pathToName(imp.importPath, p.srcDir),
+		unknown = append(unknown, imp.importPath)
+	}
+
+	if !p.useGoPackages || !p.loadRealPackageNames {
+		pathToName := importPathToNameBasic
+		if p.loadRealPackageNames {
+			pathToName = importPathToName
+		}
+		for _, path := range unknown {
+			p.knownPackages[path] = &packageInfo{
+				name:    pathToName(path, p.srcDir),
+				exports: map[string]bool{},
+			}
+		}
+		return nil
+	}
+
+	cfg := newPackagesConfig(packages.LoadFiles)
+	pkgs, err := packages.Load(cfg, unknown...)
+	if err != nil {
+		return err
+	}
+	for _, pkg := range pkgs {
+		p.knownPackages[VendorlessPath(pkg.PkgPath)] = &packageInfo{
+			name:    pkg.Name,
 			exports: map[string]bool{},
 		}
 	}
+	// We may not have found all the packages. Guess the rest.
+	for _, path := range unknown {
+		if _, ok := p.knownPackages[path]; ok {
+			continue
+		}
+		p.knownPackages[path] = &packageInfo{
+			name:    importPathToNameBasic(path, p.srcDir),
+			exports: map[string]bool{},
+		}
+	}
+	return nil
 }
 
 // importIdentifier returns the indentifier that imp will introduce.
@@ -342,7 +382,11 @@ func (p *pass) fix() bool {
 		}
 	}
 
-	if p.addImportNames {
+	for _, imp := range selected {
+		astutil.AddNamedImport(p.fset, p.f, imp.name, imp.importPath)
+	}
+
+	if p.loadRealPackageNames {
 		for _, imp := range p.f.Imports {
 			if imp.Name != nil {
 				continue
@@ -358,10 +402,6 @@ func (p *pass) fix() bool {
 		}
 	}
 
-	for _, imp := range selected {
-		astutil.AddNamedImport(p.fset, p.f, imp.name, imp.importPath)
-	}
-
 	return true
 }
 
@@ -395,6 +435,9 @@ func (p *pass) assumeSiblingImportsValid() {
 func (p *pass) addCandidate(imp *importInfo, pkg *packageInfo) {
 	p.candidates = append(p.candidates, imp)
 	if existing, ok := p.knownPackages[imp.importPath]; ok {
+		if existing.name == "" {
+			existing.name = pkg.name
+		}
 		for export := range pkg.exports {
 			existing.exports[export] = true
 		}
@@ -425,7 +468,6 @@ func fixImportsDefault(fset *token.FileSet, f *ast.File, filename string) error
 	// complete. We can't add any imports yet, because we don't know
 	// if missing references are actually package vars.
 	p := &pass{fset: fset, f: f, srcDir: srcDir}
-	p.pathToName = importPathToNameBasic
 	if p.load() {
 		return nil
 	}
@@ -435,7 +477,6 @@ func fixImportsDefault(fset *token.FileSet, f *ast.File, filename string) error
 	// Second pass: add information from other files in the same package,
 	// like their package vars and imports.
 	p = &pass{fset: fset, f: f, srcDir: srcDir}
-	p.pathToName = importPathToNameBasic
 	p.otherFiles = otherFiles
 	if p.load() {
 		return nil
@@ -448,11 +489,14 @@ func fixImportsDefault(fset *token.FileSet, f *ast.File, filename string) error
 		return nil
 	}
 
+	// The only things that use go/packages happen in the third pass,
+	// so we can delay calling go env until this point.
+	useGoPackages := shouldUseGoPackages()
+
 	// Third pass: get real package names where we had previously used
 	// the naive algorithm.
-	p = &pass{fset: fset, f: f, srcDir: srcDir}
-	p.pathToName = importPathToName
-	p.addImportNames = true
+	p = &pass{fset: fset, f: f, srcDir: srcDir, useGoPackages: useGoPackages}
+	p.loadRealPackageNames = true
 	p.otherFiles = otherFiles
 	if p.load() {
 		return nil
@@ -466,7 +510,7 @@ func fixImportsDefault(fset *token.FileSet, f *ast.File, filename string) error
 
 	// Go look for candidates in $GOPATH, etc. We don't necessarily load
 	// the real exports of sibling imports, so keep assuming their contents.
-	if err := addGoPathCandidates(p, p.missing, filename); err != nil {
+	if err := addExternalCandidates(p, p.missing, filename); err != nil {
 		return err
 	}
 
@@ -475,6 +519,37 @@ func fixImportsDefault(fset *token.FileSet, f *ast.File, filename string) error
 	return nil
 }
 
+// Values controlling the use of go/packages, for testing only.
+var forceGoPackages, _ = strconv.ParseBool(os.Getenv("GOIMPORTSFORCEGOPACKAGES"))
+var goPackagesDir string
+var go111ModuleEnv string
+
+func shouldUseGoPackages() bool {
+	if forceGoPackages {
+		return true
+	}
+
+	cmd := exec.Command("go", "env", "GOMOD")
+	cmd.Dir = goPackagesDir
+	out, err := cmd.Output()
+	if err != nil {
+		return false
+	}
+	return len(bytes.TrimSpace(out)) > 0
+}
+
+func newPackagesConfig(mode packages.LoadMode) *packages.Config {
+	cfg := &packages.Config{
+		Mode: mode,
+		Dir:  goPackagesDir,
+		Env:  append(os.Environ(), "GOROOT="+build.Default.GOROOT, "GOPATH="+build.Default.GOPATH),
+	}
+	if go111ModuleEnv != "" {
+		cfg.Env = append(cfg.Env, "GO111MODULE="+go111ModuleEnv)
+	}
+	return cfg
+}
+
 func addStdlibCandidates(pass *pass, refs map[string]map[string]bool) {
 	add := func(pkg string) {
 		pass.addCandidate(
@@ -496,7 +571,43 @@ func addStdlibCandidates(pass *pass, refs map[string]map[string]bool) {
 	}
 }
 
-func addGoPathCandidates(pass *pass, refs map[string]map[string]bool, filename string) error {
+func scanGoPackages(refs map[string]map[string]bool) ([]*pkg, error) {
+	var loadQueries []string
+	for pkgName := range refs {
+		loadQueries = append(loadQueries, "name="+pkgName)
+	}
+	sort.Strings(loadQueries)
+	cfg := newPackagesConfig(packages.LoadTypes)
+	goPackages, err := packages.Load(cfg, loadQueries...)
+	if err != nil {
+		return nil, err
+	}
+
+	var scan []*pkg
+	for _, goPackage := range goPackages {
+		scan = append(scan, &pkg{
+			dir:             filepath.Dir(goPackage.CompiledGoFiles[0]),
+			importPathShort: VendorlessPath(goPackage.PkgPath),
+			goPackage:       goPackage,
+		})
+	}
+	return scan, nil
+}
+
+var addExternalCandidates = addExternalCandidatesDefault
+
+func addExternalCandidatesDefault(pass *pass, refs map[string]map[string]bool, filename string) error {
+	var dirScan []*pkg
+	if pass.useGoPackages {
+		var err error
+		dirScan, err = scanGoPackages(refs)
+		if err != nil {
+			return err
+		}
+	} else {
+		dirScan = scanGoDirs()
+	}
+
 	// Search for imports matching potential package references.
 	type result struct {
 		imp *importInfo
@@ -519,7 +630,7 @@ func addGoPathCandidates(pass *pass, refs map[string]map[string]bool, filename s
 		go func(pkgName string, symbols map[string]bool) {
 			defer wg.Done()
 
-			ipath, err := findImport(ctx, pkgName, symbols, filename)
+			found, err := findImport(ctx, dirScan, pkgName, symbols, filename)
 
 			if err != nil {
 				firstErrOnce.Do(func() {
@@ -529,17 +640,12 @@ func addGoPathCandidates(pass *pass, refs map[string]map[string]bool, filename s
 				return
 			}
 
-			if ipath == "" {
+			if found == nil {
 				return // No matching package.
 			}
 
 			imp := &importInfo{
-				importPath: ipath,
-			}
-			// If the package name isn't what you'd expect looking
-			// at the import path, add an explicit name.
-			if importPathToNameBasic(ipath, pass.srcDir) != pkgName {
-				imp.name = pkgName
+				importPath: found.importPathShort,
 			}
 
 			pkg := &packageInfo{
@@ -649,16 +755,9 @@ func importPathToNameGoPathParse(importPath, srcDir string) (packageName string,
 	return "", fmt.Errorf("no importable package found in %d Go files", nfile)
 }
 
-// Directory-scanning state.
-var (
-	// scanOnce guards calling scanGoDirs and assigning dirScan
-	scanOnce sync.Once
-	dirScan  map[string]*pkg // abs dir path => *pkg
-)
-
 type pkg struct {
+	goPackage       *packages.Package
 	dir             string // absolute file path to pkg directory ("/usr/lib/go/src/net/http")
-	importPath      string // full pkg import path ("net/http", "foo/bar/vendor/a/b")
 	importPathShort string // vendorless import path ("net/http", "a/b")
 }
 
@@ -704,23 +803,25 @@ func distance(basepath, targetpath string) int {
 }
 
 // scanGoDirs populates the dirScan map for GOPATH and GOROOT.
-func scanGoDirs() map[string]*pkg {
-	result := make(map[string]*pkg)
+func scanGoDirs() []*pkg {
+	dupCheck := make(map[string]bool)
+	var result []*pkg
+
 	var mu sync.Mutex
 
 	add := func(root gopathwalk.Root, dir string) {
 		mu.Lock()
 		defer mu.Unlock()
 
-		if _, dup := result[dir]; dup {
+		if _, dup := dupCheck[dir]; dup {
 			return
 		}
+		dupCheck[dir] = true
 		importpath := filepath.ToSlash(dir[len(root.Path)+len("/"):])
-		result[dir] = &pkg{
-			importPath:      importpath,
+		result = append(result, &pkg{
 			importPathShort: VendorlessPath(importpath),
 			dir:             dir,
-		}
+		})
 	}
 	gopathwalk.Walk(gopathwalk.SrcDirsRoots(), add, gopathwalk.Options{Debug: Debug, ModulesEnabled: false})
 	return result
@@ -741,9 +842,19 @@ func VendorlessPath(ipath string) string {
 
 // loadExports returns the set of exported symbols in the package at dir.
 // It returns nil on error or if the package name in dir does not match expectPackage.
-func loadExports(ctx context.Context, expectPackage, dir string) (map[string]bool, error) {
+func loadExports(ctx context.Context, expectPackage string, pkg *pkg) (map[string]bool, error) {
+	if pkg.goPackage != nil {
+		exports := map[string]bool{}
+		for _, name := range pkg.goPackage.Types.Scope().Names() {
+			if ast.IsExported(name) {
+				exports[name] = true
+			}
+		}
+		return exports, nil
+	}
+
 	if Debug {
-		log.Printf("loading exports in dir %s (seeking package %s)", dir, expectPackage)
+		log.Printf("loading exports in dir %s (seeking package %s)", pkg.dir, expectPackage)
 	}
 	exports := make(map[string]bool)
 
@@ -767,7 +878,7 @@ func loadExports(ctx context.Context, expectPackage, dir string) (map[string]boo
 		return notTests, nil
 	}
 
-	files, err := buildCtx.ReadDir(dir)
+	files, err := buildCtx.ReadDir(pkg.dir)
 	if err != nil {
 		log.Print(err)
 		return nil, err
@@ -782,11 +893,11 @@ func loadExports(ctx context.Context, expectPackage, dir string) (map[string]boo
 		default:
 		}
 
-		match, err := buildCtx.MatchFile(dir, fi.Name())
+		match, err := buildCtx.MatchFile(pkg.dir, fi.Name())
 		if err != nil || !match {
 			continue
 		}
-		fullFile := filepath.Join(dir, fi.Name())
+		fullFile := filepath.Join(pkg.dir, fi.Name())
 		f, err := parser.ParseFile(fset, fullFile, nil, 0)
 		if err != nil {
 			if Debug {
@@ -801,7 +912,7 @@ func loadExports(ctx context.Context, expectPackage, dir string) (map[string]boo
 			continue
 		}
 		if pkgName != expectPackage {
-			err := fmt.Errorf("scan of dir %v is not expected package %v (actually %v)", dir, expectPackage, pkgName)
+			err := fmt.Errorf("scan of dir %v is not expected package %v (actually %v)", pkg.dir, expectPackage, pkgName)
 			if Debug {
 				log.Print(err)
 			}
@@ -820,21 +931,19 @@ func loadExports(ctx context.Context, expectPackage, dir string) (map[string]boo
 			exportList = append(exportList, k)
 		}
 		sort.Strings(exportList)
-		log.Printf("loaded exports in dir %v (package %v): %v", dir, expectPackage, strings.Join(exportList, ", "))
+		log.Printf("loaded exports in dir %v (package %v): %v", pkg.dir, expectPackage, strings.Join(exportList, ", "))
 	}
 	return exports, nil
 }
 
 // findImport searches for a package with the given symbols.
 // If no package is found, findImport returns ("", false, nil)
-func findImport(ctx context.Context, pkgName string, symbols map[string]bool, filename string) (foundPkg string, err error) {
+func findImport(ctx context.Context, dirScan []*pkg, pkgName string, symbols map[string]bool, filename string) (*pkg, error) {
 	pkgDir, err := filepath.Abs(filename)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	pkgDir = filepath.Dir(pkgDir)
-	// Scan $GOROOT and each $GOPATH.
-	scanOnce.Do(func() { dirScan = scanGoDirs() })
 
 	// Find candidate packages, looking only at their directory names first.
 	var candidates []pkgDistance
@@ -891,7 +1000,7 @@ func findImport(ctx context.Context, pkgName string, symbols map[string]bool, fi
 					wg.Done()
 				}()
 
-				exports, err := loadExports(ctx, pkgName, c.pkg.dir)
+				exports, err := loadExports(ctx, pkgName, c.pkg)
 				if err != nil {
 					resc <- nil
 					return
@@ -915,9 +1024,9 @@ func findImport(ctx context.Context, pkgName string, symbols map[string]bool, fi
 		if pkg == nil {
 			continue
 		}
-		return pkg.importPathShort, nil
+		return pkg, nil
 	}
-	return "", nil
+	return nil, nil
 }
 
 // pkgIsCandidate reports whether pkg is a candidate for satisfying the

imports/fix_test.go

@@ -10,7 +10,6 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
-	"sync"
 	"testing"
 
 	"golang.org/x/tools/go/packages/packagestest"
@@ -1449,6 +1448,7 @@ import "golang.org/x/net/http2/hpack"
 var _ = hpack.HuffmanDecode
 `
 	testConfig{
+		gopathOnly: true,
 		module: packagestest.Module{
 			Name: "foo.com",
 			Files: fm{
@@ -1498,15 +1498,34 @@ type fm map[string]interface{}
 func (c testConfig) test(t *testing.T, fn func(*goimportTest)) {
 	t.Helper()
 
-	var exported *packagestest.Exported
 	if c.module.Name != "" {
 		c.modules = []packagestest.Module{c.module}
 	}
 
-	for _, exporter := range []packagestest.Exporter{packagestest.GOPATH} {
-		t.Run(exporter.Name(), func(t *testing.T) {
+	kinds := []string{"GOPATH_GoPackages"}
+	for _, exporter := range packagestest.All {
+		kinds = append(kinds, exporter.Name())
+	}
+	for _, kind := range kinds {
+		t.Run(kind, func(t *testing.T) {
 			t.Helper()
-			exported = packagestest.Export(t, exporter, c.modules)
+
+			var exporter packagestest.Exporter
+			switch kind {
+			case "GOPATH":
+				exporter = packagestest.GOPATH
+			case "GOPATH_GoPackages":
+				exporter = packagestest.GOPATH
+				forceGoPackages = true
+			case "Modules":
+				if c.gopathOnly {
+					t.Skip("test marked GOPATH-only")
+				}
+				exporter = packagestest.Modules
+			default:
+				panic("unknown test type")
+			}
+			exported := packagestest.Export(t, exporter, c.modules)
 			defer exported.Cleanup()
 
 			env := make(map[string]string)
@@ -1519,19 +1538,22 @@ func (c testConfig) test(t *testing.T, fn func(*goimportTest)) {
 			goroot := env["GOROOT"]
 			gopath := env["GOPATH"]
 
-			scanOnce = sync.Once{}
-
 			oldGOPATH := build.Default.GOPATH
 			oldGOROOT := build.Default.GOROOT
 			oldCompiler := build.Default.Compiler
 			build.Default.GOROOT = goroot
 			build.Default.GOPATH = gopath
 			build.Default.Compiler = "gc"
+			goPackagesDir = exported.Config.Dir
+			go111ModuleEnv = env["GO111MODULE"]
 
 			defer func() {
 				build.Default.GOPATH = oldGOPATH
 				build.Default.GOROOT = oldGOROOT
 				build.Default.Compiler = oldCompiler
+				go111ModuleEnv = ""
+				goPackagesDir = ""
+				forceGoPackages = false
 			}()
 
 			it := &goimportTest{
@@ -2008,7 +2030,6 @@ func TestPkgIsCandidate(t *testing.T) {
 			pkgIdent: "client",
 			pkg: &pkg{
 				dir:             "/gopath/src/client",
-				importPath:      "client",
 				importPathShort: "client",
 			},
 			want: true,
@@ -2019,7 +2040,6 @@ func TestPkgIsCandidate(t *testing.T) {
 			pkgIdent: "zzz",
 			pkg: &pkg{
 				dir:             "/gopath/src/client",
-				importPath:      "client",
 				importPathShort: "client",
 			},
 			want: false,
@@ -2030,7 +2050,6 @@ func TestPkgIsCandidate(t *testing.T) {
 			pkgIdent: "client",
 			pkg: &pkg{
 				dir:             "/gopath/src/client/foo/foo/foo",
-				importPath:      "client/foo/foo",
 				importPathShort: "client/foo/foo",
 			},
 			want: false,
@@ -2041,7 +2060,6 @@ func TestPkgIsCandidate(t *testing.T) {
 			pkgIdent: "client",
 			pkg: &pkg{
 				dir:             "/gopath/src/foo/go-client",
-				importPath:      "foo/go-client",
 				importPathShort: "foo/go-client",
 			},
 			want: true,
@@ -2052,7 +2070,6 @@ func TestPkgIsCandidate(t *testing.T) {
 			pkgIdent: "client",
 			pkg: &pkg{
 				dir:             "/gopath/src/foo/internal/client",
-				importPath:      "foo/internal/client",
 				importPathShort: "foo/internal/client",
 			},
 			want: false,
@@ -2063,7 +2080,6 @@ func TestPkgIsCandidate(t *testing.T) {
 			pkgIdent: "client",
 			pkg: &pkg{
 				dir:             "/gopath/src/foo/internal/client",
-				importPath:      "foo/internal/client",
 				importPathShort: "foo/internal/client",
 			},
 			want: true,
@@ -2074,7 +2090,6 @@ func TestPkgIsCandidate(t *testing.T) {
 			pkgIdent: "client",
 			pkg: &pkg{
 				dir:             "/gopath/src/other/vendor/client",
-				importPath:      "other/vendor/client",
 				importPathShort: "client",
 			},
 			want: false,
@@ -2085,7 +2100,6 @@ func TestPkgIsCandidate(t *testing.T) {
 			pkgIdent: "client",
 			pkg: &pkg{
 				dir:             "/gopath/src/foo/vendor/client",
-				importPath:      "other/foo/client",
 				importPathShort: "client",
 			},
 			want: true,
@@ -2096,7 +2110,6 @@ func TestPkgIsCandidate(t *testing.T) {
 			pkgIdent: "socketio",
 			pkg: &pkg{
 				dir:             "/gopath/src/foo/socket-io",
-				importPath:      "foo/socket-io",
 				importPathShort: "foo/socket-io",
 			},
 			want: true,
@@ -2107,7 +2120,6 @@ func TestPkgIsCandidate(t *testing.T) {
 			pkgIdent: "fooprod",
 			pkg: &pkg{
 				dir:             "/gopath/src/foo/FooPROD",
-				importPath:      "foo/FooPROD",
 				importPathShort: "foo/FooPROD",
 			},
 			want: true,
@@ -2118,7 +2130,6 @@ func TestPkgIsCandidate(t *testing.T) {
 			pkgIdent: "fooprod",
 			pkg: &pkg{
 				dir:             "/gopath/src/foo/Foo-PROD",
-				importPath:      "foo/Foo-PROD",
 				importPathShort: "foo/Foo-PROD",
 			},
 			want: true,