go/packages: add roots for overlaid packages for all query types

Previously, we only added roots for contains queries. Now, we borrow
some matching logic from the go command to add roots for any new
packages that originate from overlays, no matter the query pattern.
Ideally, we won't have to borrow this logic once 1.16 is released with
native overlay support.

Change-Id: Iaa06f5ecda47820bd41ed0e42d9c2d33a0539b11
Reviewed-on: https://go-review.googlesource.com/c/tools/+/254418
Trust: Rebecca Stambler <rstambler@golang.org>
Trust: 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>
Reviewed-by: Heschi Kreinick <heschi@google.com>

3 files changed, 176 insertions, 1 deletions

diff --git a/go/packages/golist.go b/go/packages/golist.go
index b4d232be7..bc04503c1 100644
--- a/go/packages/golist.go
+++ b/go/packages/golist.go
@@ -246,6 +246,21 @@ extractQueries:
 			}
 		}
 	}
+	// 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)
+			}
+		}
+	}
 
 	sizeswg.Wait()
 	if sizeserr != nil {
@@ -753,7 +768,8 @@ func (state *golistState) getGoVersion() (string, error) {
 	return state.goVersion, state.goVersionError
 }
 
-// getPkgPath finds the package path of a directory if it's relative to a root directory.
+// getPkgPath finds the package path of a directory if it's relative to a root
+// directory.
 func (state *golistState) getPkgPath(dir string) (string, bool, error) {
 	absDir, err := filepath.Abs(dir)
 	if err != nil {
diff --git a/go/packages/golist_overlay.go b/go/packages/golist_overlay.go
index 154facc64..9ddecb27c 100644
--- a/go/packages/golist_overlay.go
+++ b/go/packages/golist_overlay.go
@@ -8,6 +8,7 @@ import (
 	"log"
 	"os"
 	"path/filepath"
+	"regexp"
 	"sort"
 	"strconv"
 	"strings"
@@ -481,3 +482,79 @@ func maybeFixPackageName(newName string, isTestFile bool, pkgsOfDir []*Package)
 		p.Name = newName
 	}
 }
+
+// This function is copy-pasted from
+// https://github.com/golang/go/blob/9706f510a5e2754595d716bd64be8375997311fb/src/cmd/go/internal/search/search.go#L360.
+// It should be deleted when we remove support for overlays from go/packages.
+//
+// NOTE: This does not handle any ./... or ./ style queries, as this function
+// doesn't know the working directory.
+//
+// matchPattern(pattern)(name) reports whether
+// name matches pattern. Pattern is a limited glob
+// pattern in which '...' means 'any string' and there
+// is no other special syntax.
+// Unfortunately, there are two special cases. Quoting "go help packages":
+//
+// First, /... at the end of the pattern can match an empty string,
+// so that net/... matches both net and packages in its subdirectories, like net/http.
+// Second, any slash-separated pattern element containing a wildcard never
+// participates in a match of the "vendor" element in the path of a vendored
+// package, so that ./... does not match packages in subdirectories of
+// ./vendor or ./mycode/vendor, but ./vendor/... and ./mycode/vendor/... do.
+// Note, however, that a directory named vendor that itself contains code
+// is not a vendored package: cmd/vendor would be a command named vendor,
+// and the pattern cmd/... matches it.
+func matchPattern(pattern string) func(name string) bool {
+	// Convert pattern to regular expression.
+	// The strategy for the trailing /... is to nest it in an explicit ? expression.
+	// The strategy for the vendor exclusion is to change the unmatchable
+	// vendor strings to a disallowed code point (vendorChar) and to use
+	// "(anything but that codepoint)*" as the implementation of the ... wildcard.
+	// This is a bit complicated but the obvious alternative,
+	// namely a hand-written search like in most shell glob matchers,
+	// is too easy to make accidentally exponential.
+	// Using package regexp guarantees linear-time matching.
+
+	const vendorChar = "\x00"
+
+	if strings.Contains(pattern, vendorChar) {
+		return func(name string) bool { return false }
+	}
+
+	re := regexp.QuoteMeta(pattern)
+	re = replaceVendor(re, vendorChar)
+	switch {
+	case strings.HasSuffix(re, `/`+vendorChar+`/\.\.\.`):
+		re = strings.TrimSuffix(re, `/`+vendorChar+`/\.\.\.`) + `(/vendor|/` + vendorChar + `/\.\.\.)`
+	case re == vendorChar+`/\.\.\.`:
+		re = `(/vendor|/` + vendorChar + `/\.\.\.)`
+	case strings.HasSuffix(re, `/\.\.\.`):
+		re = strings.TrimSuffix(re, `/\.\.\.`) + `(/\.\.\.)?`
+	}
+	re = strings.ReplaceAll(re, `\.\.\.`, `[^`+vendorChar+`]*`)
+
+	reg := regexp.MustCompile(`^` + re + `$`)
+
+	return func(name string) bool {
+		if strings.Contains(name, vendorChar) {
+			return false
+		}
+		return reg.MatchString(replaceVendor(name, vendorChar))
+	}
+}
+
+// replaceVendor returns the result of replacing
+// non-trailing vendor path elements in x with repl.
+func replaceVendor(x, repl string) string {
+	if !strings.Contains(x, "vendor") {
+		return x
+	}
+	elem := strings.Split(x, "/")
+	for i := 0; i < len(elem)-1; i++ {
+		if elem[i] == "vendor" {
+			elem[i] = repl
+		}
+	}
+	return strings.Join(elem, "/")
+}
diff --git a/go/packages/overlay_test.go b/go/packages/overlay_test.go
index b70a875ec..c77251e7d 100644
--- a/go/packages/overlay_test.go
+++ b/go/packages/overlay_test.go
@@ -932,3 +932,85 @@ func _() {
 		t.Fatalf(`expected import "os", found none: %v`, pkg.Imports)
 	}
 }
+
+// Tests that overlays are applied for different kinds of load patterns.
+func TestLoadDifferentPatterns(t *testing.T) {
+	packagestest.TestAll(t, testLoadDifferentPatterns)
+}
+func testLoadDifferentPatterns(t *testing.T, exporter packagestest.Exporter) {
+	exported := packagestest.Export(t, exporter, []packagestest.Module{
+		{
+			Name: "golang.org/fake",
+			Files: map[string]interface{}{
+				"foo.txt": "placeholder",
+				"b/b.go": `package b
+import "golang.org/fake/a"
+func _() {
+	a.Hi()
+}
+`,
+			},
+		},
+	})
+	defer exported.Cleanup()
+
+	exported.Config.Mode = everythingMode
+	exported.Config.Tests = true
+
+	dir := filepath.Dir(exported.File("golang.org/fake", "foo.txt"))
+	exported.Config.Overlay = map[string][]byte{
+		filepath.Join(dir, "a", "a.go"): []byte(`package a
+import "fmt"
+func Hi() {
+	fmt.Println("")
+}
+`),
+	}
+	for _, tc := range []struct {
+		pattern string
+	}{
+		{"golang.org/fake/a"},
+		{"golang.org/fake/..."},
+		{fmt.Sprintf("file=%s", filepath.Join(dir, "a", "a.go"))},
+	} {
+		t.Run(tc.pattern, func(t *testing.T) {
+			initial, err := packages.Load(exported.Config, tc.pattern)
+			if err != nil {
+				t.Fatal(err)
+			}
+			var match *packages.Package
+			for _, pkg := range initial {
+				if pkg.PkgPath == "golang.org/fake/a" {
+					match = pkg
+					break
+				}
+			}
+			if match == nil {
+				t.Fatalf(`expected package path "golang.org/fake/a", got %q`, match.PkgPath)
+			}
+			if match.PkgPath != "golang.org/fake/a" {
+				t.Fatalf(`expected package path "golang.org/fake/a", got %q`, match.PkgPath)
+			}
+			if _, ok := match.Imports["fmt"]; !ok {
+				t.Fatalf(`expected import "fmt", got none`)
+			}
+		})
+	}
+
+	// Now, load "golang.org/fake/b" and confirm that "golang.org/fake/a" is
+	// not returned as a root.
+	initial, err := packages.Load(exported.Config, "golang.org/fake/b")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(initial) > 1 {
+		t.Fatalf("expected 1 package, got %v", initial)
+	}
+	pkg := initial[0]
+	if pkg.PkgPath != "golang.org/fake/b" {
+		t.Fatalf(`expected package path "golang.org/fake/b", got %q`, pkg.PkgPath)
+	}
+	if _, ok := pkg.Imports["golang.org/fake/a"]; !ok {
+		t.Fatalf(`expected import "golang.org/fake/a", got none`)
+	}
+}