go/packages: support new main packages in overlays

Refactor the overlay code to create package structs for
new packages that don't already exist. This requires
calling out to the go command to determine module
roots to figure out which module a package belongs to.
The extra go list call is done in sequence in this CL
but can easily be done in parallel with other go list
calls in the future.

Change-Id: Ia0f7812fba250d154033038cb1e2afa7dedf0e16
Reviewed-on: https://go-review.googlesource.com/c/tools/+/179600
Run-TryBot: Michael Matloob <matloob@golang.org>
Reviewed-by: Ian Cottrell <iancottrell@google.com>

go/packages/golist_overlay.go

@@ -1,11 +1,15 @@
 package packages
 
 import (
+	"bytes"
+	"encoding/json"
 	"go/parser"
 	"go/token"
+	"path"
 	"path/filepath"
 	"strconv"
 	"strings"
+	"sync"
 )
 
 // processGolistOverlay provides rudimentary support for adding
@@ -27,52 +31,123 @@ func processGolistOverlay(cfg *Config, response *driverResponse) (modifiedPkgs,
 		havePkgs[pkg.PkgPath] = pkg.ID
 	}
 
-outer:
-	for path, contents := range cfg.Overlay {
-		base := filepath.Base(path)
-		if strings.HasSuffix(path, "_test.go") {
+	var rootDirs map[string]string
+	var onceGetRootDirs sync.Once
+
+	for opath, contents := range cfg.Overlay {
+		base := filepath.Base(opath)
+		if strings.HasSuffix(opath, "_test.go") {
 			// Overlays don't support adding new test files yet.
 			// TODO(matloob): support adding new test files.
 			continue
 		}
-		dir := filepath.Dir(path)
-		for _, pkg := range response.Packages {
-			var dirContains, fileExists bool
-			for _, f := range pkg.GoFiles {
-				if sameFile(filepath.Dir(f), dir) {
-					dirContains = true
+		dir := filepath.Dir(opath)
+		var pkg *Package
+		var fileExists bool
+		for _, p := range response.Packages {
+			for _, f := range p.GoFiles {
+				if !sameFile(filepath.Dir(f), dir) {
+					continue
 				}
+				pkg = p
 				if filepath.Base(f) == base {
 					fileExists = true
 				}
 			}
-			// The overlay could have included an entirely new package.
-			isNewPackage := extractPackage(pkg, path, contents)
-			if dirContains || isNewPackage {
-				if !fileExists {
-					pkg.GoFiles = append(pkg.GoFiles, path) // TODO(matloob): should the file just be added to GoFiles?
-					pkg.CompiledGoFiles = append(pkg.CompiledGoFiles, path)
-					modifiedPkgsSet[pkg.ID] = true
-				}
-				imports, err := extractImports(path, contents)
+		}
+		// The overlay could have included an entirely new package.
+		if pkg == nil {
+			onceGetRootDirs.Do(func() {
+				rootDirs = determineRootDirs(cfg)
+			})
+			// Try to find the module or gopath dir the file is contained in.
+			// Then for modules, add the module opath to the beginning.
+			var pkgPath string
+			for rdir, rpath := range rootDirs {
+				// TODO(matloob): This doesn't properly handle symlinks.
+				r, err := filepath.Rel(rdir, dir)
 				if err != nil {
-					// Let the parser or type checker report errors later.
-					continue outer
+					continue
 				}
-				for _, imp := range imports {
-					_, found := pkg.Imports[imp]
-					if !found {
-						needPkgsSet[imp] = true
-						// TODO(matloob): Handle cases when the following block isn't correct.
-						// These include imports of test variants, imports of vendored packages, etc.
-						id, ok := havePkgs[imp]
-						if !ok {
-							id = imp
-						}
-						pkg.Imports[imp] = &Package{ID: id}
-					}
+				pkgPath = filepath.ToSlash(r)
+				if rpath != "" {
+					pkgPath = path.Join(rpath, pkgPath)
 				}
-				continue outer
+				// We only create one new package even it can belong in multiple modules or GOPATH entries.
+				// This is okay because tools (such as the LSP) that use overlays will recompute the overlay
+				// once the file is saved, and golist will do the right thing.
+				// TODO(matloob): Implement module tiebreaking?
+				break
+			}
+			if pkgPath == "" {
+				continue
+			}
+			pkgName, ok := extractPackageName(opath, contents)
+			if !ok {
+				continue
+			}
+			id := pkgPath
+			// Try to reclaim a package with the same id if it exists in the response.
+			for _, p := range response.Packages {
+				if reclaimPackage(p, id, opath, contents) {
+					pkg = p
+					break
+				}
+			}
+			// Otherwise, create a new package
+			if pkg == nil {
+				pkg = &Package{PkgPath: pkgPath, ID: id, Name: pkgName, Imports: make(map[string]*Package)}
+				// TODO(matloob): Is it okay to amend response.Packages this way?
+				response.Packages = append(response.Packages, pkg)
+				havePkgs[pkg.PkgPath] = id
+			}
+		}
+		if !fileExists {
+			pkg.GoFiles = append(pkg.GoFiles, opath)
+			// TODO(matloob): Adding the file to CompiledGoFiles can exhibit the wrong behavior
+			// if the file will be ignored due to its build tags.
+			pkg.CompiledGoFiles = append(pkg.CompiledGoFiles, opath)
+			modifiedPkgsSet[pkg.ID] = true
+		}
+		imports, err := extractImports(opath, contents)
+		if err != nil {
+			// Let the parser or type checker report errors later.
+			continue
+		}
+		for _, imp := range imports {
+			_, found := pkg.Imports[imp]
+			if !found {
+				// TODO(matloob): Handle cases when the following block isn't correct.
+				// These include imports of test variants, imports of vendored packages, etc.
+				id, ok := havePkgs[imp]
+				if !ok {
+					id = imp
+				}
+				pkg.Imports[imp] = &Package{ID: id}
+			}
+		}
+		continue
+	}
+
+	// toPkgPath tries to guess the package path given the id.
+	// This isn't always correct -- it's certainly wrong for
+	// vendored packages' paths.
+	toPkgPath := func(id string) string {
+		// TODO(matloob): Handle vendor paths.
+		i := strings.IndexByte(id, ' ')
+		if i >= 0 {
+			return id[:i]
+		}
+		return id
+	}
+
+	// Do another pass now that new packages have been created to determine the
+	// set of missing packages.
+	for _, pkg := range response.Packages {
+		for _, imp := range pkg.Imports {
+			pkgPath := toPkgPath(imp.ID)
+			if _, ok := havePkgs[pkgPath]; !ok {
+				needPkgsSet[pkgPath] = true
 			}
 		}
 	}
@@ -88,6 +163,46 @@ outer:
 	return modifiedPkgs, needPkgs, err
 }
 
+// determineRootDirs returns a mapping from directories code can be contained in to the
+// corresponding import path prefixes of those directories.
+// Its result is used to try to determine the import path for a package containing
+// an overlay file.
+func determineRootDirs(cfg *Config) map[string]string {
+	// Assume modules first:
+	out, err := invokeGo(cfg, "list", "-m", "-json", "all")
+	if err != nil {
+		return determineRootDirsGOPATH(cfg)
+	}
+	m := map[string]string{}
+	type jsonMod struct{ Path, Dir string }
+	for dec := json.NewDecoder(out); dec.More(); {
+		mod := new(jsonMod)
+		if err := dec.Decode(mod); err != nil {
+			return m // Give up and return an empty map. Package won't be found for overlay.
+		}
+		if mod.Dir != "" && mod.Path != "" {
+			// This is a valid module; add it to the map.
+			m[mod.Dir] = mod.Path
+		}
+	}
+	return m
+}
+
+func determineRootDirsGOPATH(cfg *Config) map[string]string {
+	m := map[string]string{}
+	out, err := invokeGo(cfg, "env", "GOPATH")
+	if err != nil {
+		// Could not determine root dir mapping. Everything is best-effort, so just return an empty map.
+		// When we try to find the import path for a directory, there will be no root-dir match and
+		// we'll give up.
+		return m
+	}
+	for _, p := range filepath.SplitList(string(bytes.TrimSpace(out.Bytes()))) {
+		m[filepath.Join(p, "src")] = ""
+	}
+	return m
+}
+
 func extractImports(filename string, contents []byte) ([]string, error) {
 	f, err := parser.ParseFile(token.NewFileSet(), filename, contents, parser.ImportsOnly) // TODO(matloob): reuse fileset?
 	if err != nil {
@@ -105,13 +220,16 @@ func extractImports(filename string, contents []byte) ([]string, error) {
 	return res, nil
 }
 
-// extractPackage attempts to extract a package defined in an overlay.
+// reclaimPackage attempts to reuse a package that failed to load in an overlay.
 //
 // If the package has errors and has no Name, GoFiles, or Imports,
 // then it's possible that it doesn't yet exist on disk.
-func extractPackage(pkg *Package, filename string, contents []byte) bool {
+func reclaimPackage(pkg *Package, id string, filename string, contents []byte) bool {
 	// TODO(rstambler): Check the message of the actual error?
 	// It differs between $GOPATH and module mode.
+	if pkg.ID != id {
+		return false
+	}
 	if len(pkg.Errors) != 1 {
 		return false
 	}
@@ -124,15 +242,21 @@ func extractPackage(pkg *Package, filename string, contents []byte) bool {
 	if len(pkg.Imports) > 0 {
 		return false
 	}
-	f, err := parser.ParseFile(token.NewFileSet(), filename, contents, parser.PackageClauseOnly) // TODO(matloob): reuse fileset?
-	if err != nil {
+	pkgName, ok := extractPackageName(filename, contents)
+	if !ok {
 		return false
 	}
-	// TODO(rstambler): This doesn't work for main packages.
-	if filepath.Base(pkg.PkgPath) != f.Name.Name {
-		return false
-	}
-	pkg.Name = f.Name.Name
+	pkg.Name = pkgName
 	pkg.Errors = nil
 	return true
 }
+
+func extractPackageName(filename string, contents []byte) (string, bool) {
+	// TODO(rstambler): Check the message of the actual error?
+	// It differs between $GOPATH and module mode.
+	f, err := parser.ParseFile(token.NewFileSet(), filename, contents, parser.PackageClauseOnly) // TODO(matloob): reuse fileset?
+	if err != nil {
+		return "", false
+	}
+	return f.Name.Name, true
+}

go/packages/packages_test.go

@@ -979,6 +979,10 @@ func testNewPackagesInOverlay(t *testing.T, exporter packagestest.Exporter) {
 			filepath.Join(dir, "h", "h.go"):      []byte(`package h; const H = "h"`),
 		},
 			`"efgh_"`},
+		// Overlay with package main.
+		{map[string][]byte{
+			filepath.Join(dir, "e", "main.go"): []byte(`package main; import "golang.org/fake/a"; const E = "e" + a.A; func main(){}`)},
+			`"eabc"`},
 	} {
 		exported.Config.Overlay = test.overlay
 		exported.Config.Mode = packages.LoadAllSyntax