go/packages: use native overlay support for 1.16

This change modifies go/packages to use the go command's -overlay flag
if used with Go 1.16. It does so by adding a new Overlay field to the
gocommand.Invocation struct. go/packages writes out the overlay files
as expected by go list before invoking a `go list` command.

Fixes golang/go#41598

Change-Id: Iec5edf19ce2936d5a633d076905622c2cf779bcc
Reviewed-on: https://go-review.googlesource.com/c/tools/+/263984
Trust: Rebecca Stambler <rstambler@golang.org>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>
diff --git a/go/packages/golist.go b/go/packages/golist.go
index 787783c..81381fa 100644
--- a/go/packages/golist.go
+++ b/go/packages/golist.go
@@ -10,6 +10,7 @@
 	"encoding/json"
 	"fmt"
 	"go/types"
+	"io/ioutil"
 	"log"
 	"os"
 	"os/exec"
@@ -208,56 +209,58 @@
 		}
 	}
 
-	modifiedPkgs, needPkgs, err := state.processGolistOverlay(response)
-	if err != nil {
-		return nil, err
-	}
+	// Only use go/packages' overlay processing if we're using a Go version
+	// below 1.16. Otherwise, go list handles it.
+	if goVersion, err := state.getGoVersion(); err == nil && goVersion < 16 {
+		modifiedPkgs, needPkgs, err := state.processGolistOverlay(response)
+		if err != nil {
+			return nil, err
+		}
 
-	var containsCandidates []string
-	if len(containFiles) > 0 {
-		containsCandidates = append(containsCandidates, modifiedPkgs...)
-		containsCandidates = append(containsCandidates, needPkgs...)
-	}
-	if err := state.addNeededOverlayPackages(response, needPkgs); err != nil {
-		return nil, err
-	}
-	// Check candidate packages for containFiles.
-	if len(containFiles) > 0 {
-		for _, id := range containsCandidates {
-			pkg, ok := response.seenPackages[id]
-			if !ok {
-				response.addPackage(&Package{
-					ID: id,
-					Errors: []Error{
-						{
+		var containsCandidates []string
+		if len(containFiles) > 0 {
+			containsCandidates = append(containsCandidates, modifiedPkgs...)
+			containsCandidates = append(containsCandidates, needPkgs...)
+		}
+		if err := state.addNeededOverlayPackages(response, needPkgs); err != nil {
+			return nil, err
+		}
+		// Check candidate packages for containFiles.
+		if len(containFiles) > 0 {
+			for _, id := range containsCandidates {
+				pkg, ok := response.seenPackages[id]
+				if !ok {
+					response.addPackage(&Package{
+						ID: id,
+						Errors: []Error{{
 							Kind: ListError,
 							Msg:  fmt.Sprintf("package %s expected but not seen", id),
-						},
-					},
-				})
-				continue
-			}
-			for _, f := range containFiles {
-				for _, g := range pkg.GoFiles {
-					if sameFile(f, g) {
-						response.addRoot(id)
+						}},
+					})
+					continue
+				}
+				for _, f := range containFiles {
+					for _, g := range pkg.GoFiles {
+						if sameFile(f, g) {
+							response.addRoot(id)
+						}
 					}
 				}
 			}
 		}
-	}
-	// Add root for any package that matches a pattern. This applies only to
-	// packages that are modified by overlays, since they are not added as
-	// roots automatically.
-	for _, pattern := range restPatterns {
-		match := matchPattern(pattern)
-		for _, pkgID := range modifiedPkgs {
-			pkg, ok := response.seenPackages[pkgID]
-			if !ok {
-				continue
-			}
-			if match(pkg.PkgPath) {
-				response.addRoot(pkg.ID)
+		// Add root for any package that matches a pattern. This applies only to
+		// packages that are modified by overlays, since they are not added as
+		// roots automatically.
+		for _, pattern := range restPatterns {
+			match := matchPattern(pattern)
+			for _, pkgID := range modifiedPkgs {
+				pkg, ok := response.seenPackages[pkgID]
+				if !ok {
+					continue
+				}
+				if match(pkg.PkgPath) {
+					response.addRoot(pkg.ID)
+				}
 			}
 		}
 	}
@@ -835,6 +838,26 @@
 	cfg := state.cfg
 
 	inv := state.cfgInvocation()
+
+	// For Go versions 1.16 and above, `go list` accepts overlays directly via
+	// the -overlay flag. Set it, if it's available.
+	//
+	// The check for "list" is not necessarily required, but we should avoid
+	// getting the go version if possible.
+	if verb == "list" {
+		goVersion, err := state.getGoVersion()
+		if err != nil {
+			return nil, err
+		}
+		if goVersion >= 16 {
+			filename, cleanup, err := state.writeOverlays()
+			if err != nil {
+				return nil, err
+			}
+			defer cleanup()
+			inv.Overlay = filename
+		}
+	}
 	inv.Verb = verb
 	inv.Args = args
 	gocmdRunner := cfg.gocmdRunner
@@ -976,6 +999,67 @@
 	return stdout, nil
 }
 
+// OverlayJSON is the format overlay files are expected to be in.
+// The Replace map maps from overlaid paths to replacement paths:
+// the Go command will forward all reads trying to open
+// each overlaid path to its replacement path, or consider the overlaid
+// path not to exist if the replacement path is empty.
+//
+// From golang/go#39958.
+type OverlayJSON struct {
+	Replace map[string]string `json:"replace,omitempty"`
+}
+
+// writeOverlays writes out files for go list's -overlay flag, as described
+// above.
+func (state *golistState) writeOverlays() (filename string, cleanup func(), err error) {
+	// Do nothing if there are no overlays in the config.
+	if len(state.cfg.Overlay) == 0 {
+		return "", func() {}, nil
+	}
+	dir, err := ioutil.TempDir("", "gopackages-*")
+	if err != nil {
+		return "", nil, err
+	}
+	// The caller must clean up this directory, unless this function returns an
+	// error.
+	cleanup = func() {
+		os.RemoveAll(dir)
+	}
+	defer func() {
+		if err != nil {
+			cleanup()
+		}
+	}()
+	overlays := map[string]string{}
+	for k, v := range state.cfg.Overlay {
+		// Create a unique filename for the overlaid files, to avoid
+		// creating nested directories.
+		noSeparator := strings.Join(strings.Split(filepath.ToSlash(k), "/"), "")
+		f, err := ioutil.TempFile(dir, fmt.Sprintf("*-%s", noSeparator))
+		if err != nil {
+			return "", func() {}, err
+		}
+		if _, err := f.Write(v); err != nil {
+			return "", func() {}, err
+		}
+		if err := f.Close(); err != nil {
+			return "", func() {}, err
+		}
+		overlays[k] = f.Name()
+	}
+	b, err := json.Marshal(OverlayJSON{Replace: overlays})
+	if err != nil {
+		return "", func() {}, err
+	}
+	// Write out the overlay file that contains the filepath mappings.
+	filename = filepath.Join(dir, "overlay.json")
+	if err := ioutil.WriteFile(filename, b, 0665); err != nil {
+		return "", func() {}, err
+	}
+	return filename, cleanup, nil
+}
+
 func containsGoFile(s []string) bool {
 	for _, f := range s {
 		if strings.HasSuffix(f, ".go") {
diff --git a/go/packages/overlay_test.go b/go/packages/overlay_test.go
index a706570..8ce5a7c 100644
--- a/go/packages/overlay_test.go
+++ b/go/packages/overlay_test.go
@@ -87,6 +87,9 @@
 		{"fake [fake.test]", "foox", 2},
 		{"fake.test", "main", 1},
 	}
+	if len(initial) != 3 {
+		t.Fatalf("expected 3 packages, got %v", len(initial))
+	}
 	for i := 0; i < 3; i++ {
 		if ok := checkPkg(t, initial[i], want[i].id, want[i].name, want[i].count); !ok {
 			t.Errorf("%d: got {%s %s %d}, expected %v", i, initial[i].ID,
@@ -102,7 +105,8 @@
 	packagestest.TestAll(t, testOverlayChangesTestPackageName)
 }
 func testOverlayChangesTestPackageName(t *testing.T, exporter packagestest.Exporter) {
-	log.SetFlags(log.Lshortfile)
+	testenv.NeedsGo1Point(t, 16)
+
 	exported := packagestest.Export(t, exporter, []packagestest.Module{{
 		Name: "fake",
 		Files: map[string]interface{}{
@@ -127,10 +131,13 @@
 		id, name string
 		count    int
 	}{
-		{"fake", "foo", 0},
+		{"fake", "foox", 0},
 		{"fake [fake.test]", "foox", 1},
 		{"fake.test", "main", 1},
 	}
+	if len(initial) != 3 {
+		t.Fatalf("expected 3 packages, got %v", len(initial))
+	}
 	for i := 0; i < 3; i++ {
 		if ok := checkPkg(t, initial[i], want[i].id, want[i].name, want[i].count); !ok {
 			t.Errorf("got {%s %s %d}, expected %v", initial[i].ID,
@@ -329,6 +336,9 @@
 
 	// Find package golang.org/fake/c
 	sort.Slice(pkgs, func(i, j int) bool { return pkgs[i].ID < pkgs[j].ID })
+	if len(pkgs) != 2 {
+		t.Fatalf("expected 2 packages, got %v", len(pkgs))
+	}
 	pkgc := pkgs[0]
 	if pkgc.ID != "golang.org/fake/c" {
 		t.Errorf("expected first package in sorted list to be \"golang.org/fake/c\", got %v", pkgc.ID)
@@ -804,6 +814,9 @@
 				if err != nil {
 					t.Fatal(err)
 				}
+				if len(initial) != 1 {
+					t.Fatalf("expected 1 packages, got %v", len(initial))
+				}
 				pkg := initial[0]
 				if pkg.ID != tt.wantID {
 					t.Fatalf("expected package ID %q, got %q", tt.wantID, pkg.ID)
@@ -986,7 +999,7 @@
 				}
 			}
 			if match == nil {
-				t.Fatalf(`expected package path "golang.org/fake/a", got %q`, match.PkgPath)
+				t.Fatalf(`expected package path "golang.org/fake/a", got none`)
 			}
 			if match.PkgPath != "golang.org/fake/a" {
 				t.Fatalf(`expected package path "golang.org/fake/a", got %q`, match.PkgPath)
@@ -1072,6 +1085,9 @@
 	if err != nil {
 		t.Error(err)
 	}
+	if len(initial) != 1 {
+		t.Fatalf(`expected 1 package, got %v`, len(initial))
+	}
 	pkg := initial[0]
 	if pkg.PkgPath != "b.com/inner" {
 		t.Fatalf(`expected package path "b.com/inner", got %q`, pkg.PkgPath)
diff --git a/internal/gocommand/invoke.go b/internal/gocommand/invoke.go
index b5c061b..8ba2253 100644
--- a/internal/gocommand/invoke.go
+++ b/internal/gocommand/invoke.go
@@ -132,6 +132,7 @@
 	BuildFlags []string
 	ModFlag    string
 	ModFile    string
+	Overlay    string
 	Env        []string
 	WorkingDir string
 	Logf       func(format string, args ...interface{})
@@ -171,6 +172,11 @@
 			goArgs = append(goArgs, "-mod="+i.ModFlag)
 		}
 	}
+	appendOverlayFlag := func() {
+		if i.Overlay != "" {
+			goArgs = append(goArgs, "-overlay="+i.Overlay)
+		}
+	}
 
 	switch i.Verb {
 	case "env", "version":
@@ -189,6 +195,7 @@
 		goArgs = append(goArgs, i.BuildFlags...)
 		appendModFile()
 		appendModFlag()
+		appendOverlayFlag()
 		goArgs = append(goArgs, i.Args...)
 	}
 	cmd := exec.Command("go", goArgs...)