go/packages: support overlays for replaced modules

This change implements the approach described in
https://github.com/golang/go/issues/37629#issuecomment-594179751.
This feature is necessary for multi-module workspace support in gopls.

Updates golang/go#32394

Change-Id: Iab328020a2681f8651405922cc25d9d06cb1ac82
Reviewed-on: https://go-review.googlesource.com/c/tools/+/254368
Trust: Rebecca Stambler <rstambler@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
diff --git a/go/packages/golist_overlay.go b/go/packages/golist_overlay.go
index 9ddecb2..874f901 100644
--- a/go/packages/golist_overlay.go
+++ b/go/packages/golist_overlay.go
@@ -12,6 +12,8 @@
 	"sort"
 	"strconv"
 	"strings"
+
+	"golang.org/x/tools/internal/gocommand"
 )
 
 // processGolistOverlay provides rudimentary support for adding
@@ -326,24 +328,25 @@
 }
 
 func (state *golistState) determineRootDirsModules() (map[string]string, error) {
-	// This will only return the root directory for the main module.
-	// For now we only support overlays in main modules.
+	// List all of the modules--the first will be the directory for the main
+	// module. Any replaced modules will also need to be treated as roots.
 	// Editing files in the module cache isn't a great idea, so we don't
-	// plan to ever support that, but editing files in replaced modules
-	// is something we may want to support. To do that, we'll want to
-	// do a go list -m to determine the replaced module's module path and
-	// directory, and then a go list -m {{with .Replace}}{{.Dir}}{{end}} <replaced module's path>
-	// from the main module to determine if that module is actually a replacement.
-	// See bcmills's comment here: https://github.com/golang/go/issues/37629#issuecomment-594179751
-	// for more information.
-	out, err := state.invokeGo("list", "-m", "-json")
+	// plan to ever support that.
+	out, err := state.invokeGo("list", "-m", "-json", "all")
 	if err != nil {
-		return nil, err
+		// 'go list all' will fail if we're outside of a module and
+		// GO111MODULE=on. Try falling back without 'all'.
+		var innerErr error
+		out, innerErr = state.invokeGo("list", "-m", "-json")
+		if innerErr != nil {
+			return nil, err
+		}
 	}
-	m := map[string]string{}
-	type jsonMod struct{ Path, Dir string }
+	roots := map[string]string{}
+	modules := map[string]string{}
+	var i int
 	for dec := json.NewDecoder(out); dec.More(); {
-		mod := new(jsonMod)
+		mod := new(gocommand.ModuleJSON)
 		if err := dec.Decode(mod); err != nil {
 			return nil, err
 		}
@@ -353,10 +356,15 @@
 			if err != nil {
 				return nil, err
 			}
-			m[absDir] = mod.Path
+			modules[absDir] = mod.Path
+			// The first result is the main module.
+			if i == 0 || mod.Replace != nil && mod.Replace.Path != "" {
+				roots[absDir] = mod.Path
+			}
 		}
+		i++
 	}
-	return m, nil
+	return roots, nil
 }
 
 func (state *golistState) determineRootDirsGOPATH() (map[string]string, error) {
diff --git a/go/packages/overlay_test.go b/go/packages/overlay_test.go
index c77251e..a706570 100644
--- a/go/packages/overlay_test.go
+++ b/go/packages/overlay_test.go
@@ -1014,3 +1014,69 @@
 		t.Fatalf(`expected import "golang.org/fake/a", got none`)
 	}
 }
+
+// Tests that overlays are applied for a replaced module.
+// This does not use go/packagestest because it needs to write a replace
+// directive with an absolute path in one of the module's go.mod files.
+func TestOverlaysInReplace(t *testing.T) {
+	// Create module b.com in a temporary directory. Do not add any Go files
+	// on disk.
+	tmpPkgs, err := ioutil.TempDir("", "modules")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tmpPkgs)
+
+	dirB := filepath.Join(tmpPkgs, "b")
+	if err := os.Mkdir(dirB, 0775); err != nil {
+		t.Fatal(err)
+	}
+	if err := ioutil.WriteFile(filepath.Join(dirB, "go.mod"), []byte(fmt.Sprintf("module %s.com", dirB)), 0775); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.MkdirAll(filepath.Join(dirB, "inner"), 0775); err != nil {
+		t.Fatal(err)
+	}
+
+	// Create a separate module that requires and replaces b.com.
+	tmpWorkspace, err := ioutil.TempDir("", "workspace")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tmpWorkspace)
+	goModContent := fmt.Sprintf(`module workspace.com
+
+require (
+	b.com v0.0.0-00010101000000-000000000000
+)
+
+replace (
+	b.com => %s
+)
+`, dirB)
+	if err := ioutil.WriteFile(filepath.Join(tmpWorkspace, "go.mod"), []byte(goModContent), 0775); err != nil {
+		t.Fatal(err)
+	}
+
+	// Add Go files for b.com/inner in an overlay and try loading it from the
+	// workspace.com module.
+	config := &packages.Config{
+		Dir:  tmpWorkspace,
+		Mode: packages.LoadAllSyntax,
+		Logf: t.Logf,
+		Overlay: map[string][]byte{
+			filepath.Join(dirB, "inner", "b.go"): []byte(`package inner; import "fmt"; func _() { fmt.Println("");`),
+		},
+	}
+	initial, err := packages.Load(config, "b.com/...")
+	if err != nil {
+		t.Error(err)
+	}
+	pkg := initial[0]
+	if pkg.PkgPath != "b.com/inner" {
+		t.Fatalf(`expected package path "b.com/inner", got %q`, pkg.PkgPath)
+	}
+	if _, ok := pkg.Imports["fmt"]; !ok {
+		t.Fatalf(`expected import "fmt", got none`)
+	}
+}