go/packages: add support to overlays for unwritten files to existing packages

This support is not perfect, but should produce more accurate results than
the previous behavior, which is to silently drop those files.

Change-Id: Ia00c042bf303e9ec0fb1cbd579c0fccb29073de0
Reviewed-on: https://go-review.googlesource.com/c/151999
Run-TryBot: Michael Matloob <matloob@golang.org>
Reviewed-by: Ian Cottrell <iancottrell@google.com>
diff --git a/go/packages/golist.go b/go/packages/golist.go
index 6be07b9..5181f6d 100644
--- a/go/packages/golist.go
+++ b/go/packages/golist.go
@@ -118,10 +118,6 @@
 	// types.SizesFor always returns nil or a *types.StdSizes
 	response.Sizes, _ = sizes.(*types.StdSizes)
 
-	if len(containFiles) == 0 && len(packagesNamed) == 0 {
-		return response, nil
-	}
-
 	seenPkgs := make(map[string]*Package) // for deduplication. different containing queries could produce same packages
 	for _, pkg := range response.Packages {
 		seenPkgs[pkg.ID] = pkg
@@ -134,20 +130,44 @@
 		response.Packages = append(response.Packages, p)
 	}
 
-	containsResults, err := runContainsQueries(cfg, listfunc, isFallback, addPkg, containFiles)
-	if err != nil {
-		return nil, err
+	if len(containFiles) != 0 {
+		containsResults, err := runContainsQueries(cfg, listfunc, isFallback, addPkg, containFiles)
+		if err != nil {
+			return nil, err
+		}
+		response.Roots = append(response.Roots, containsResults...)
 	}
-	response.Roots = append(response.Roots, containsResults...)
 
-	namedResults, err := runNamedQueries(cfg, listfunc, addPkg, packagesNamed)
+	if len(packagesNamed) != 0 {
+		namedResults, err := runNamedQueries(cfg, listfunc, addPkg, packagesNamed)
+		if err != nil {
+			return nil, err
+		}
+		response.Roots = append(response.Roots, namedResults...)
+	}
+
+	needPkgs, err := processGolistOverlay(cfg, response)
 	if err != nil {
 		return nil, err
 	}
-	response.Roots = append(response.Roots, namedResults...)
+	if len(needPkgs) > 0 {
+		addNeededOverlayPackages(cfg, listfunc, addPkg, needPkgs)
+	}
+
 	return response, nil
 }
 
+func addNeededOverlayPackages(cfg *Config, driver driver, addPkg func(*Package), pkgs []string) error {
+	response, err := driver(cfg, pkgs...)
+	if err != nil {
+		return err
+	}
+	for _, pkg := range response.Packages {
+		addPkg(pkg)
+	}
+	return nil
+}
+
 func runContainsQueries(cfg *Config, driver driver, isFallback bool, addPkg func(*Package), queries []string) ([]string, error) {
 	var results []string
 	for _, query := range queries {
diff --git a/go/packages/golist_overlay.go b/go/packages/golist_overlay.go
new file mode 100644
index 0000000..60438a1
--- /dev/null
+++ b/go/packages/golist_overlay.go
@@ -0,0 +1,98 @@
+package packages
+
+import (
+	"go/parser"
+	"go/token"
+	"path/filepath"
+	"strconv"
+	"strings"
+)
+
+// processGolistOverlay provides rudimentary support for adding
+// files that don't exist on disk to an overlay. The results can be
+// sometimes incorrect.
+// TODO(matloob): Handle unsupported cases, including the following:
+// - test files
+// - adding test and non-test files to test variants of packages
+// - determining the correct package to add given a new import path
+// - creating packages that don't exist
+func processGolistOverlay(cfg *Config, response *driverResponse) (needPkgs []string, err error) {
+	havePkgs := make(map[string]string) // importPath -> non-test package ID
+	needPkgsSet := make(map[string]bool)
+
+	for _, pkg := range response.Packages {
+		// This is an approximation of import path to id. This can be
+		// wrong for tests, vendored packages, and a number of other cases.
+		havePkgs[pkg.PkgPath] = pkg.ID
+	}
+
+outer:
+	for path, contents := range cfg.Overlay {
+		base := filepath.Base(path)
+		if strings.HasSuffix(path, "_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
+				}
+				if filepath.Base(f) == base {
+					fileExists = true
+				}
+			}
+			if dirContains {
+				if !fileExists {
+					pkg.GoFiles = append(pkg.GoFiles, path) // TODO(matloob): should the file just be added to GoFiles?
+					pkg.CompiledGoFiles = append(pkg.CompiledGoFiles, path)
+				}
+				imports, err := extractImports(path, contents)
+				if err != nil {
+					// Let the parser or type checker report errors later.
+					continue outer
+				}
+				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}
+					}
+				}
+				continue outer
+			}
+		}
+	}
+
+	needPkgs = make([]string, 0, len(needPkgsSet))
+	for pkg := range needPkgsSet {
+		needPkgs = append(needPkgs, pkg)
+	}
+	return needPkgs, err
+}
+
+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 {
+		return nil, err
+	}
+	var res []string
+	for _, imp := range f.Imports {
+		quotedPath := imp.Path.Value
+		path, err := strconv.Unquote(quotedPath)
+		if err != nil {
+			return nil, err
+		}
+		res = append(res, path)
+	}
+	return res, nil
+}
diff --git a/go/packages/packages.go b/go/packages/packages.go
index c1110f3..81e5b82 100644
--- a/go/packages/packages.go
+++ b/go/packages/packages.go
@@ -432,6 +432,7 @@
 				ld.Mode >= LoadTypes && rootIndex >= 0,
 			needsrc: ld.Mode >= LoadAllSyntax ||
 				ld.Mode >= LoadSyntax && rootIndex >= 0 ||
+				ld.Overlay != nil || // Overlays can invalidate export data. TODO(matloob): make this check fine-grained based on dependencies on overlaid files
 				pkg.ExportFile == "" && pkg.PkgPath != "unsafe",
 		}
 		ld.pkgs[lpkg.ID] = lpkg
@@ -819,6 +820,15 @@
 // the same file.
 //
 func sameFile(x, y string) bool {
+	if x == y {
+		// It could be the case that y doesn't exist.
+		// For instance, it may be an overlay file that
+		// hasn't been written to disk. To handle that case
+		// let x == y through. (We added the exact absolute path
+		// string to the CompiledGoFiles list, so the unwritten
+		// overlay case implies x==y.)
+		return true
+	}
 	if filepath.Base(x) == filepath.Base(y) { // (optimisation)
 		if xi, err := os.Stat(x); err == nil {
 			if yi, err := os.Stat(y); err == nil {
diff --git a/go/packages/packages_test.go b/go/packages/packages_test.go
index 50e6344..2548e57 100644
--- a/go/packages/packages_test.go
+++ b/go/packages/packages_test.go
@@ -824,8 +824,18 @@
 		{map[string][]byte{}, `"abc"`, nil}, // empty overlay
 		{map[string][]byte{exported.File("golang.org/fake", "c/c.go"): []byte(`package c; const C = "C"`)}, `"abC"`, nil},
 		{map[string][]byte{exported.File("golang.org/fake", "b/b.go"): []byte(`package b; import "golang.org/fake/c"; const B = "B" + c.C`)}, `"aBc"`, nil},
-		{map[string][]byte{exported.File("golang.org/fake", "b/b.go"): []byte(`package b; import "d"; const B = "B" + d.D`)}, `unknown`,
-			[]string{`could not import d (no metadata for d)`}},
+		// Overlay with an existing file in an existing package adding a new import.
+		{map[string][]byte{exported.File("golang.org/fake", "b/b.go"): []byte(`package b; import "golang.org/fake/d"; const B = "B" + d.D`)}, `"aBd"`, nil},
+		// Overlay with a new file in an existing package.
+		{map[string][]byte{
+			exported.File("golang.org/fake", "c/c.go"):                                               []byte(`package c;`),
+			filepath.Join(filepath.Dir(exported.File("golang.org/fake", "c/c.go")), "c_new_file.go"): []byte(`package c; const C = "Ç"`)},
+			`"abÇ"`, nil},
+		// Overlay with a new file in an existing package, adding a new dependency to that package.
+		{map[string][]byte{
+			exported.File("golang.org/fake", "c/c.go"):                                               []byte(`package c;`),
+			filepath.Join(filepath.Dir(exported.File("golang.org/fake", "c/c.go")), "c_new_file.go"): []byte(`package c; import "golang.org/fake/d"; const C = "c" + d.D`)},
+			`"abcd"`, nil},
 	} {
 		exported.Config.Overlay = test.overlay
 		exported.Config.Mode = packages.LoadAllSyntax