gopls/internal/lsp/source: find references in test packages

When finding references or implementations, we must include objects
defined in intermediate test variants. However, as we have seen we
should be careful to avoid accidentally type-checking intermediate test
variants in ParseFull, which is not their workspace parse mode.
Therefore eliminate the problematic TypecheckAll type-check mode in
favor of special handling in this one case where it is necessary.

Along the way:
 - Simplify the mapping of protocol position->offset. This should not
   require getting a package, or even parsing a file. For isolation,
   just use the NewColumnMapper constructor, even though it may
   technically result in building a token.File multiple times.
 - Update package renaming logic to use TypecheckWorkspace, since we
   only need to rename within the workspace.
 - Add regtest support for Implementations requests.

Fixes golang/go#43144

Change-Id: I41f684ad766f5af805abbd7c5ee0a010ff9b9b8c
Reviewed-on: https://go-review.googlesource.com/c/tools/+/438755
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Run-TryBot: Robert Findley <rfindley@google.com>
Reviewed-by: Alan Donovan <adonovan@google.com>

gopls/internal/lsp/cache/load.go

@@ -655,6 +655,7 @@ func computeWorkspacePackagesLocked(s *snapshot, meta *metadataGraph) map[Packag
 		if !m.Valid {
 			continue
 		}
+
 		if !containsPackageLocked(s, m.Metadata) {
 			continue
 		}

gopls/internal/lsp/cache/snapshot.go

@@ -698,27 +698,16 @@ func (s *snapshot) packageHandlesForFile(ctx context.Context, uri span.URI, mode
 		if m := s.getMetadata(id); m != nil && m.IsIntermediateTestVariant() && !withIntermediateTestVariants {
 			continue
 		}
-		var parseModes []source.ParseMode
-		switch mode {
-		case source.TypecheckAll:
-			if s.workspaceParseMode(id) == source.ParseFull {
-				parseModes = []source.ParseMode{source.ParseFull}
-			} else {
-				parseModes = []source.ParseMode{source.ParseExported, source.ParseFull}
-			}
-		case source.TypecheckFull:
-			parseModes = []source.ParseMode{source.ParseFull}
-		case source.TypecheckWorkspace:
-			parseModes = []source.ParseMode{s.workspaceParseMode(id)}
+		parseMode := source.ParseFull
+		if mode == source.TypecheckWorkspace {
+			parseMode = s.workspaceParseMode(id)
 		}
 
-		for _, parseMode := range parseModes {
-			ph, err := s.buildPackageHandle(ctx, id, parseMode)
-			if err != nil {
-				return nil, err
-			}
-			phs = append(phs, ph)
+		ph, err := s.buildPackageHandle(ctx, id, parseMode)
+		if err != nil {
+			return nil, err
 		}
+		phs = append(phs, ph)
 	}
 	return phs, nil
 }

gopls/internal/lsp/fake/editor.go

@@ -1137,7 +1137,8 @@ func (e *Editor) InlayHint(ctx context.Context, path string) ([]protocol.InlayHi
 	return hints, nil
 }
 
-// References executes a reference request on the server.
+// References returns references to the object at (path, pos), as returned by
+// the connected LSP server. If no server is connected, it returns (nil, nil).
 func (e *Editor) References(ctx context.Context, path string, pos Pos) ([]protocol.Location, error) {
 	if e.Server == nil {
 		return nil, nil
@@ -1164,6 +1165,8 @@ func (e *Editor) References(ctx context.Context, path string, pos Pos) ([]protoc
 	return locations, nil
 }
 
+// Rename performs a rename of the object at (path, pos) to newName, using the
+// connected LSP server. If no server is connected, it returns nil.
 func (e *Editor) Rename(ctx context.Context, path string, pos Pos, newName string) error {
 	if e.Server == nil {
 		return nil
@@ -1185,6 +1188,28 @@ func (e *Editor) Rename(ctx context.Context, path string, pos Pos, newName strin
 	return nil
 }
 
+// Implementations returns implementations for the object at (path, pos), as
+// returned by the connected LSP server. If no server is connected, it returns
+// (nil, nil).
+func (e *Editor) Implementations(ctx context.Context, path string, pos Pos) ([]protocol.Location, error) {
+	if e.Server == nil {
+		return nil, nil
+	}
+	e.mu.Lock()
+	_, ok := e.buffers[path]
+	e.mu.Unlock()
+	if !ok {
+		return nil, fmt.Errorf("buffer %q is not open", path)
+	}
+	params := &protocol.ImplementationParams{
+		TextDocumentPositionParams: protocol.TextDocumentPositionParams{
+			TextDocument: e.TextDocumentIdentifier(path),
+			Position:     pos.ToProtocolPosition(),
+		},
+	}
+	return e.Server.Implementation(ctx, params)
+}
+
 func (e *Editor) RenameFile(ctx context.Context, oldPath, newPath string) error {
 	closed, opened, err := e.renameBuffers(ctx, oldPath, newPath)
 	if err != nil {

gopls/internal/lsp/regtest/wrappers.go

@@ -389,8 +389,7 @@ func (e *Env) WorkspaceSymbol(sym string) []protocol.SymbolInformation {
 	return ans
 }
 
-// References calls textDocument/references for the given path at the given
-// position.
+// References wraps Editor.References, calling t.Fatal on any error.
 func (e *Env) References(path string, pos fake.Pos) []protocol.Location {
 	e.T.Helper()
 	locations, err := e.Editor.References(e.Ctx, path, pos)
@@ -408,6 +407,16 @@ func (e *Env) Rename(path string, pos fake.Pos, newName string) {
 	}
 }
 
+// Implementations wraps Editor.Implementations, calling t.Fatal on any error.
+func (e *Env) Implementations(path string, pos fake.Pos) []protocol.Location {
+	e.T.Helper()
+	locations, err := e.Editor.Implementations(e.Ctx, path, pos)
+	if err != nil {
+		e.T.Fatal(err)
+	}
+	return locations
+}
+
 // RenameFile wraps Editor.RenameFile, calling t.Fatal on any error.
 func (e *Env) RenameFile(oldPath, newPath string) {
 	e.T.Helper()

gopls/internal/lsp/source/implementation.go

@@ -208,36 +208,31 @@ var (
 	errNoObjectFound = errors.New("no object found")
 )
 
-// qualifiedObjsAtProtocolPos returns info for all the type.Objects
-// referenced at the given position. An object will be returned for
-// every package that the file belongs to, in every typechecking mode
-// applicable.
+// qualifiedObjsAtProtocolPos returns info for all the types.Objects referenced
+// at the given position, for the following selection of packages:
+//
+// 1. all packages (including all test variants), in their workspace parse mode
+// 2. if not included above, at least one package containing uri in full parse mode
+//
+// Finding objects in (1) ensures that we locate references within all
+// workspace packages, including in x_test packages. Including (2) ensures that
+// we find local references in the current package, for non-workspace packages
+// that may be open.
 func qualifiedObjsAtProtocolPos(ctx context.Context, s Snapshot, uri span.URI, pp protocol.Position) ([]qualifiedObject, error) {
-	offset, err := protocolPositionToOffset(ctx, s, uri, pp)
+	fh, err := s.GetFile(ctx, uri)
 	if err != nil {
 		return nil, err
 	}
-	return qualifiedObjsAtLocation(ctx, s, positionKey{uri, offset}, map[positionKey]bool{})
-}
-
-func protocolPositionToOffset(ctx context.Context, s Snapshot, uri span.URI, pp protocol.Position) (int, error) {
-	pkgs, err := s.PackagesForFile(ctx, uri, TypecheckAll, false)
+	content, err := fh.Read()
 	if err != nil {
-		return 0, err
-	}
-	if len(pkgs) == 0 {
-		return 0, errNoObjectFound
-	}
-	pkg := pkgs[0]
-	pgf, err := pkg.File(uri)
-	if err != nil {
-		return 0, err
+		return nil, err
 	}
-	pos, err := pgf.Mapper.Pos(pp)
+	m := protocol.NewColumnMapper(uri, content)
+	offset, err := m.Offset(pp)
 	if err != nil {
-		return 0, err
+		return nil, err
 	}
-	return safetoken.Offset(pgf.Tok, pos)
+	return qualifiedObjsAtLocation(ctx, s, positionKey{uri, offset}, map[positionKey]bool{})
 }
 
 // A positionKey identifies a byte offset within a file (URI).
@@ -251,8 +246,8 @@ type positionKey struct {
 	offset int
 }
 
-// qualifiedObjsAtLocation finds all objects referenced at offset in uri, across
-// all packages in the snapshot.
+// qualifiedObjsAtLocation finds all objects referenced at offset in uri,
+// across all packages in the snapshot.
 func qualifiedObjsAtLocation(ctx context.Context, s Snapshot, key positionKey, seen map[positionKey]bool) ([]qualifiedObject, error) {
 	if seen[key] {
 		return nil, nil
@@ -268,11 +263,31 @@ func qualifiedObjsAtLocation(ctx context.Context, s Snapshot, key positionKey, s
 	// try to be comprehensive in case we ever support variations on build
 	// constraints.
 
-	pkgs, err := s.PackagesForFile(ctx, key.uri, TypecheckAll, false)
+	pkgs, err := s.PackagesForFile(ctx, key.uri, TypecheckWorkspace, true)
 	if err != nil {
 		return nil, err
 	}
 
+	// In order to allow basic references/rename/implementations to function when
+	// non-workspace packages are open, ensure that we have at least one fully
+	// parsed package for the current file. This allows us to find references
+	// inside the open package. Use WidestPackage to capture references in test
+	// files.
+	hasFullPackage := false
+	for _, pkg := range pkgs {
+		if pkg.ParseMode() == ParseFull {
+			hasFullPackage = true
+			break
+		}
+	}
+	if !hasFullPackage {
+		pkg, err := s.PackageForFile(ctx, key.uri, TypecheckFull, WidestPackage)
+		if err != nil {
+			return nil, err
+		}
+		pkgs = append(pkgs, pkg)
+	}
+
 	// report objects in the order we encounter them. This ensures that the first
 	// result is at the cursor...
 	var qualifiedObjs []qualifiedObject

gopls/internal/lsp/source/references.go

@@ -62,7 +62,7 @@ func References(ctx context.Context, s Snapshot, f FileHandle, pp protocol.Posit
 	}
 
 	if inPackageName {
-		renamingPkg, err := s.PackageForFile(ctx, f.URI(), TypecheckAll, NarrowestPackage)
+		renamingPkg, err := s.PackageForFile(ctx, f.URI(), TypecheckWorkspace, NarrowestPackage)
 		if err != nil {
 			return nil, err
 		}

gopls/internal/lsp/source/rename.go

@@ -76,7 +76,7 @@ func PrepareRename(ctx context.Context, snapshot Snapshot, f FileHandle, pp prot
 			err := errors.New("can't rename package: LSP client does not support file renaming")
 			return nil, err, err
 		}
-		renamingPkg, err := snapshot.PackageForFile(ctx, f.URI(), TypecheckAll, NarrowestPackage)
+		renamingPkg, err := snapshot.PackageForFile(ctx, f.URI(), TypecheckWorkspace, NarrowestPackage)
 		if err != nil {
 			return nil, err, err
 		}
@@ -169,7 +169,7 @@ func Rename(ctx context.Context, s Snapshot, f FileHandle, pp protocol.Position,
 		// TODO(rfindley): but is this correct? What about x_test packages that
 		// import the renaming package?
 		const includeTestVariants = false
-		pkgs, err := s.PackagesForFile(ctx, f.URI(), TypecheckAll, includeTestVariants)
+		pkgs, err := s.PackagesForFile(ctx, f.URI(), TypecheckWorkspace, includeTestVariants)
 		if err != nil {
 			return nil, true, err
 		}

gopls/internal/lsp/source/view.go

@@ -437,16 +437,11 @@ var AllParseModes = []ParseMode{ParseHeader, ParseExported, ParseFull}
 type TypecheckMode int
 
 const (
-	// Invalid default value.
-	TypecheckUnknown TypecheckMode = iota
 	// TypecheckFull means to use ParseFull.
-	TypecheckFull
+	TypecheckFull TypecheckMode = iota
 	// TypecheckWorkspace means to use ParseFull for workspace packages, and
 	// ParseExported for others.
 	TypecheckWorkspace
-	// TypecheckAll means ParseFull for workspace packages, and both Full and
-	// Exported for others. Only valid for some functions.
-	TypecheckAll
 )
 
 type VersionedFileHandle interface {

gopls/internal/regtest/misc/references_test.go

@@ -6,9 +6,12 @@ package misc
 
 import (
 	"fmt"
+	"sort"
 	"strings"
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
+	"golang.org/x/tools/gopls/internal/lsp/protocol"
 	. "golang.org/x/tools/gopls/internal/lsp/regtest"
 )
 
@@ -169,3 +172,119 @@ func main() {
 		}
 	})
 }
+
+// Test for golang/go#43144.
+//
+// Verify that we search for references and implementations in intermediate
+// test variants.
+func TestReferencesInTestVariants(t *testing.T) {
+	const files = `
+-- go.mod --
+module foo.mod
+
+go 1.12
+-- foo/foo.go --
+package foo
+
+import "foo.mod/bar"
+
+const Foo = 42
+
+type T int
+type Interface interface{ M() }
+
+func _() {
+	_ = bar.Blah
+}
+
+-- bar/bar.go --
+package bar
+
+var Blah = 123
+
+-- bar/bar_test.go --
+package bar
+
+type Mer struct{}
+func (Mer) M() {}
+
+func TestBar() {
+	_ = Blah
+}
+-- bar/bar_x_test.go --
+package bar_test
+
+import (
+	"foo.mod/bar"
+	"foo.mod/foo"
+)
+
+type Mer struct{}
+func (Mer) M() {}
+
+func _() {
+	_ = bar.Blah
+	_ = foo.Foo
+}
+`
+
+	Run(t, files, func(t *testing.T, env *Env) {
+		env.OpenFile("foo/foo.go")
+
+		// Helper to map locations relative file paths.
+		fileLocations := func(locs []protocol.Location) []string {
+			var got []string
+			for _, loc := range locs {
+				got = append(got, env.Sandbox.Workdir.URIToPath(loc.URI))
+			}
+			sort.Strings(got)
+			return got
+		}
+
+		refTests := []struct {
+			re       string
+			wantRefs []string
+		}{
+			// Blah is referenced:
+			// - inside the foo.mod/bar (ordinary) package
+			// - inside the foo.mod/bar [foo.mod/bar.test] test variant package
+			// - from the foo.mod/bar_test [foo.mod/bar.test] x_test package
+			// - from the foo.mod/foo package
+			{"Blah", []string{"bar/bar.go", "bar/bar_test.go", "bar/bar_x_test.go", "foo/foo.go"}},
+
+			// Foo is referenced in bar_x_test.go via the intermediate test variant
+			// foo.mod/foo [foo.mod/bar.test].
+			{"Foo", []string{"bar/bar_x_test.go", "foo/foo.go"}},
+		}
+
+		for _, test := range refTests {
+			pos := env.RegexpSearch("foo/foo.go", test.re)
+			refs := env.References("foo/foo.go", pos)
+
+			got := fileLocations(refs)
+			if diff := cmp.Diff(test.wantRefs, got); diff != "" {
+				t.Errorf("References(%q) returned unexpected diff (-want +got):\n%s", test.re, diff)
+			}
+		}
+
+		implTests := []struct {
+			re        string
+			wantImpls []string
+		}{
+			// Interface is implemented both in foo.mod/bar [foo.mod/bar.test] (which
+			// doesn't import foo), and in foo.mod/bar_test [foo.mod/bar.test], which
+			// imports the test variant of foo.
+			{"Interface", []string{"bar/bar_test.go", "bar/bar_x_test.go"}},
+		}
+
+		for _, test := range implTests {
+			pos := env.RegexpSearch("foo/foo.go", test.re)
+			refs := env.Implementations("foo/foo.go", pos)
+
+			got := fileLocations(refs)
+			if diff := cmp.Diff(test.wantImpls, got); diff != "" {
+				t.Errorf("Implementations(%q) returned unexpected diff (-want +got):\n%s", test.re, diff)
+			}
+		}
+	})
+}