From d94536333cf7bae4426610a0f57c571b99aa7ccd Mon Sep 17 00:00:00 2001
From: Heschi Kreinick <heschi@google.com>
Date: Mon, 17 Aug 2020 14:36:26 -0400
Subject: [PATCH] internal/lsp/cache: don't always type check in default mode

CL 248380 forced all type checking to be in the default workspace mode.
In that CL, I said I couldn't think of any features that would break. It
appears I didn't think very hard. Navigation features inside of
dependencies are something I use all the time and they broke.

Reintroduce the ability to get packages in a particular mode, and make
it convenient to get them in all relevant modes. Update some critical
features to do so, and add regression tests.

Fixes golang/go#40809.

Change-Id: I96279f4ff994203694629ea872795246c410b206
Reviewed-on: https://go-review.googlesource.com/c/tools/+/249120
Run-TryBot: Heschi Kreinick <heschi@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>
---
 internal/lsp/cache/check.go             | 16 +++++++----
 internal/lsp/cache/load.go              |  2 +-
 internal/lsp/cache/snapshot.go          | 36 +++++++++++++++++-------
 internal/lsp/code_action.go             |  2 +-
 internal/lsp/command.go                 |  2 +-
 internal/lsp/link.go                    |  3 +-
 internal/lsp/regtest/definition_test.go | 32 +++++++++++++++------
 internal/lsp/regtest/references_test.go |  2 +-
 internal/lsp/regtest/workspace_test.go  |  5 ++--
 internal/lsp/regtest/wrappers.go        |  3 +-
 internal/lsp/source/identifier.go       | 37 +++++++++++++++++++------
 internal/lsp/source/implementation.go   |  5 ++--
 internal/lsp/source/util.go             |  2 +-
 internal/lsp/source/view.go             | 36 ++++++++++++++++++------
 14 files changed, 129 insertions(+), 54 deletions(-)

diff --git a/internal/lsp/cache/check.go b/internal/lsp/cache/check.go
index 3727eadcd61..75642c20c67 100644
--- a/internal/lsp/cache/check.go
+++ b/internal/lsp/cache/check.go
@@ -56,11 +56,7 @@ type packageData struct {
 }
 
 // buildPackageHandle returns a packageHandle for a given package and mode.
-func (s *snapshot) buildPackageHandle(ctx context.Context, id packageID) (*packageHandle, error) {
-	mode := source.ParseExported
-	if _, ok := s.isWorkspacePackage(id); ok {
-		mode = source.ParseFull
-	}
+func (s *snapshot) buildPackageHandle(ctx context.Context, id packageID, mode source.ParseMode) (*packageHandle, error) {
 	if ph := s.getPackage(id, mode); ph != nil {
 		return ph, nil
 	}
@@ -145,7 +141,7 @@ func (s *snapshot) buildKey(ctx context.Context, id packageID, mode source.Parse
 	// Begin computing the key by getting the depKeys for all dependencies.
 	var depKeys []packageHandleKey
 	for _, depID := range depList {
-		depHandle, err := s.buildPackageHandle(ctx, depID)
+		depHandle, err := s.buildPackageHandle(ctx, depID, s.workspaceParseMode(depID))
 		if err != nil {
 			event.Error(ctx, "no dep handle", err, tag.Package.Of(string(depID)))
 			if ctx.Err() != nil {
@@ -163,6 +159,14 @@ func (s *snapshot) buildKey(ctx context.Context, id packageID, mode source.Parse
 	return ph, deps, nil
 }
 
+func (s *snapshot) workspaceParseMode(id packageID) source.ParseMode {
+	if _, ws := s.isWorkspacePackage(id); ws {
+		return source.ParseFull
+	} else {
+		return source.ParseExported
+	}
+}
+
 func checkPackageKey(ctx context.Context, id packageID, pghs []*parseGoHandle, cfg *packages.Config, deps []packageHandleKey, mode source.ParseMode) packageHandleKey {
 	b := bytes.NewBuffer(nil)
 	b.WriteString(string(id))
diff --git a/internal/lsp/cache/load.go b/internal/lsp/cache/load.go
index 2e82e6ea4cf..135e2e5b92a 100644
--- a/internal/lsp/cache/load.go
+++ b/internal/lsp/cache/load.go
@@ -152,7 +152,7 @@ func (s *snapshot) load(ctx context.Context, scopes ...interface{}) error {
 		if err != nil {
 			return err
 		}
-		if _, err := s.buildPackageHandle(ctx, m.id); err != nil {
+		if _, err := s.buildPackageHandle(ctx, m.id, s.workspaceParseMode(m.id)); err != nil {
 			return err
 		}
 	}
diff --git a/internal/lsp/cache/snapshot.go b/internal/lsp/cache/snapshot.go
index 3500bff5f0d..bf9d37356b2 100644
--- a/internal/lsp/cache/snapshot.go
+++ b/internal/lsp/cache/snapshot.go
@@ -255,7 +255,7 @@ func hashUnsavedOverlays(files map[span.URI]source.VersionedFileHandle) string {
 	return hashContents([]byte(strings.Join(unsaved, "")))
 }
 
-func (s *snapshot) PackagesForFile(ctx context.Context, uri span.URI) ([]source.Package, error) {
+func (s *snapshot) PackagesForFile(ctx context.Context, uri span.URI, mode source.TypecheckMode) ([]source.Package, error) {
 	ctx = event.Label(ctx, tag.URI.Of(uri))
 
 	// Check if we should reload metadata for the file. We don't invalidate IDs
@@ -283,17 +283,33 @@ func (s *snapshot) PackagesForFile(ctx context.Context, uri span.URI) ([]source.
 	// Get the list of IDs from the snapshot again, in case it has changed.
 	var pkgs []source.Package
 	for _, id := range s.getIDsForURI(uri) {
-		pkg, err := s.checkedPackage(ctx, id)
-		if err != nil {
-			return nil, err
+		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)}
+		}
+
+		for _, parseMode := range parseModes {
+			pkg, err := s.checkedPackage(ctx, id, parseMode)
+			if err != nil {
+				return nil, err
+			}
+			pkgs = append(pkgs, pkg)
 		}
-		pkgs = append(pkgs, pkg)
 	}
 	return pkgs, nil
 }
 
-func (s *snapshot) checkedPackage(ctx context.Context, id packageID) (*pkg, error) {
-	ph, err := s.buildPackageHandle(ctx, id)
+func (s *snapshot) checkedPackage(ctx context.Context, id packageID, mode source.ParseMode) (*pkg, error) {
+	ph, err := s.buildPackageHandle(ctx, id, mode)
 	if err != nil {
 		return nil, err
 	}
@@ -312,7 +328,7 @@ func (s *snapshot) GetReverseDependencies(ctx context.Context, id string) ([]sou
 
 	var pkgs []source.Package
 	for id := range ids {
-		pkg, err := s.checkedPackage(ctx, id)
+		pkg, err := s.checkedPackage(ctx, id, s.workspaceParseMode(id))
 		if err != nil {
 			return nil, err
 		}
@@ -448,7 +464,7 @@ func (s *snapshot) WorkspacePackages(ctx context.Context) ([]source.Package, err
 	}
 	var pkgs []source.Package
 	for _, pkgID := range s.workspacePackageIDs() {
-		pkg, err := s.checkedPackage(ctx, pkgID)
+		pkg, err := s.checkedPackage(ctx, pkgID, s.workspaceParseMode(pkgID))
 		if err != nil {
 			return nil, err
 		}
@@ -476,7 +492,7 @@ func (s *snapshot) KnownPackages(ctx context.Context) ([]source.Package, error)
 
 	var pkgs []source.Package
 	for _, id := range ids {
-		pkg, err := s.checkedPackage(ctx, id)
+		pkg, err := s.checkedPackage(ctx, id, s.workspaceParseMode(id))
 		if err != nil {
 			return nil, err
 		}
diff --git a/internal/lsp/code_action.go b/internal/lsp/code_action.go
index 736c3d2fbbf..d9b90f85729 100644
--- a/internal/lsp/code_action.go
+++ b/internal/lsp/code_action.go
@@ -119,7 +119,7 @@ func (s *Server) codeAction(ctx context.Context, params *protocol.CodeActionPara
 		if ctx.Err() != nil {
 			return nil, ctx.Err()
 		}
-		pkgs, err := snapshot.PackagesForFile(ctx, fh.URI())
+		pkgs, err := snapshot.PackagesForFile(ctx, fh.URI(), source.TypecheckFull)
 		if err != nil {
 			return nil, err
 		}
diff --git a/internal/lsp/command.go b/internal/lsp/command.go
index c8cb44864a5..1e4c1268f40 100644
--- a/internal/lsp/command.go
+++ b/internal/lsp/command.go
@@ -197,7 +197,7 @@ func (s *Server) runTests(ctx context.Context, snapshot source.Snapshot, uri pro
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	pkgs, err := snapshot.PackagesForFile(ctx, uri.SpanURI())
+	pkgs, err := snapshot.PackagesForFile(ctx, uri.SpanURI(), source.TypecheckWorkspace)
 	if err != nil {
 		return err
 	}
diff --git a/internal/lsp/link.go b/internal/lsp/link.go
index 822a6cdaa20..6c0b561fb6b 100644
--- a/internal/lsp/link.go
+++ b/internal/lsp/link.go
@@ -99,7 +99,8 @@ func modLinks(ctx context.Context, snapshot source.Snapshot, fh source.FileHandl
 
 func goLinks(ctx context.Context, snapshot source.Snapshot, fh source.FileHandle) ([]protocol.DocumentLink, error) {
 	view := snapshot.View()
-	pkgs, err := snapshot.PackagesForFile(ctx, fh.URI())
+	// We don't actually need type information, so any typecheck mode is fine.
+	pkgs, err := snapshot.PackagesForFile(ctx, fh.URI(), source.TypecheckWorkspace)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/lsp/regtest/definition_test.go b/internal/lsp/regtest/definition_test.go
index 83c47f87702..727b3d57966 100644
--- a/internal/lsp/regtest/definition_test.go
+++ b/internal/lsp/regtest/definition_test.go
@@ -6,6 +6,7 @@ package regtest
 
 import (
 	"path"
+	"strings"
 	"testing"
 )
 
@@ -45,24 +46,20 @@ const stdlibDefinition = `
 -- go.mod --
 module mod.com
 
-go 1.12
 -- main.go --
 package main
 
-import (
-	"fmt"
-	"time"
-)
+import "fmt"
 
 func main() {
-	fmt.Println(time.Now())
+	fmt.Printf()
 }`
 
 func TestGoToStdlibDefinition_Issue37045(t *testing.T) {
 	runner.Run(t, stdlibDefinition, func(t *testing.T, env *Env) {
 		env.OpenFile("main.go")
-		name, pos := env.GoToDefinition("main.go", env.RegexpSearch("main.go", "Now"))
-		if got, want := path.Base(name), "time.go"; got != want {
+		name, pos := env.GoToDefinition("main.go", env.RegexpSearch("main.go", `fmt.(Printf)`))
+		if got, want := path.Base(name), "print.go"; got != want {
 			t.Errorf("GoToDefinition: got file %q, want %q", name, want)
 		}
 
@@ -77,3 +74,22 @@ func TestGoToStdlibDefinition_Issue37045(t *testing.T) {
 		}
 	})
 }
+
+func TestUnexportedStdlib_Issue40809(t *testing.T) {
+	runner.Run(t, stdlibDefinition, func(t *testing.T, env *Env) {
+		env.OpenFile("main.go")
+		name, _ := env.GoToDefinition("main.go", env.RegexpSearch("main.go", `fmt.(Printf)`))
+		env.OpenFile(name)
+
+		name, pos := env.GoToDefinition(name, env.RegexpSearch(name, `(newPrinter)\(\)`))
+		content, _ := env.Hover(name, pos)
+		if !strings.Contains(content.Value, "newPrinter") {
+			t.Fatal("definition of newPrinter went to the incorrect place")
+		}
+
+		refs := env.References(name, pos)
+		if len(refs) < 5 {
+			t.Errorf("expected 5+ references to newPrinter, found: %#v", refs)
+		}
+	})
+}
diff --git a/internal/lsp/regtest/references_test.go b/internal/lsp/regtest/references_test.go
index 2ef577ecc0f..a864092ab1f 100644
--- a/internal/lsp/regtest/references_test.go
+++ b/internal/lsp/regtest/references_test.go
@@ -8,7 +8,7 @@ import (
 	"testing"
 )
 
-func TestNonWorkspaceReferences(t *testing.T) {
+func TestStdlibReferences(t *testing.T) {
 	const files = `
 -- go.mod --
 module mod.com
diff --git a/internal/lsp/regtest/workspace_test.go b/internal/lsp/regtest/workspace_test.go
index 987a4c3cb8b..4bb31e364e4 100644
--- a/internal/lsp/regtest/workspace_test.go
+++ b/internal/lsp/regtest/workspace_test.go
@@ -106,8 +106,9 @@ func TestReferences(t *testing.T) {
 				opts = append(opts, WithRootPath(tt.rootPath))
 			}
 			withOptions(opts...).run(t, workspaceModule, func(t *testing.T, env *Env) {
-				env.OpenFile("pkg/inner/inner.go")
-				locations := env.ReferencesAtRegexp("pkg/inner/inner.go", "SaySomething")
+				f := "pkg/inner/inner.go"
+				env.OpenFile(f)
+				locations := env.References(f, env.RegexpSearch(f, `SaySomething`))
 				want := 3
 				if got := len(locations); got != want {
 					t.Fatalf("expected %v locations, got %v", want, got)
diff --git a/internal/lsp/regtest/wrappers.go b/internal/lsp/regtest/wrappers.go
index 9963160c8a8..89bc656a01d 100644
--- a/internal/lsp/regtest/wrappers.go
+++ b/internal/lsp/regtest/wrappers.go
@@ -247,9 +247,8 @@ func (e *Env) CodeLens(path string) []protocol.CodeLens {
 
 // ReferencesAtRegexp calls textDocument/references for the given path at the
 // position of the given regexp.
-func (e *Env) ReferencesAtRegexp(path string, re string) []protocol.Location {
+func (e *Env) References(path string, pos fake.Pos) []protocol.Location {
 	e.T.Helper()
-	pos := e.RegexpSearch(path, re)
 	locations, err := e.Editor.References(e.Ctx, path, pos)
 	if err != nil {
 		e.T.Fatal(err)
diff --git a/internal/lsp/source/identifier.go b/internal/lsp/source/identifier.go
index a9d5c44003f..3f4d1b099f5 100644
--- a/internal/lsp/source/identifier.go
+++ b/internal/lsp/source/identifier.go
@@ -10,6 +10,7 @@ import (
 	"go/ast"
 	"go/token"
 	"go/types"
+	"sort"
 	"strconv"
 
 	"golang.org/x/tools/internal/event"
@@ -57,19 +58,37 @@ func Identifier(ctx context.Context, snapshot Snapshot, fh FileHandle, pos proto
 	ctx, done := event.Start(ctx, "source.Identifier")
 	defer done()
 
-	pkg, pgf, err := getParsedFile(ctx, snapshot, fh, NarrowestPackage)
-	if err != nil {
-		return nil, fmt.Errorf("getting file for Identifier: %w", err)
-	}
-	spn, err := pgf.Mapper.PointSpan(pos)
+	pkgs, err := snapshot.PackagesForFile(ctx, fh.URI(), TypecheckAll)
 	if err != nil {
 		return nil, err
 	}
-	rng, err := spn.Range(pgf.Mapper.Converter)
-	if err != nil {
-		return nil, err
+	if len(pkgs) == 0 {
+		return nil, fmt.Errorf("no packages for file %v", fh.URI())
+	}
+	sort.Slice(pkgs, func(i, j int) bool {
+		return len(pkgs[i].CompiledGoFiles()) < len(pkgs[j].CompiledGoFiles())
+	})
+	var findErr error
+	for _, pkg := range pkgs {
+		pgf, err := pkg.File(fh.URI())
+		if err != nil {
+			return nil, err
+		}
+		spn, err := pgf.Mapper.PointSpan(pos)
+		if err != nil {
+			return nil, err
+		}
+		rng, err := spn.Range(pgf.Mapper.Converter)
+		if err != nil {
+			return nil, err
+		}
+		var ident *IdentifierInfo
+		ident, findErr = findIdentifier(ctx, snapshot, pkg, pgf.File, rng.Start)
+		if findErr == nil {
+			return ident, nil
+		}
 	}
-	return findIdentifier(ctx, snapshot, pkg, pgf.File, rng.Start)
+	return nil, findErr
 }
 
 var ErrNoIdentFound = errors.New("no identifier found")
diff --git a/internal/lsp/source/implementation.go b/internal/lsp/source/implementation.go
index d66b7b04588..5a59a499eb6 100644
--- a/internal/lsp/source/implementation.go
+++ b/internal/lsp/source/implementation.go
@@ -199,9 +199,10 @@ var errBuiltin = errors.New("builtin object")
 
 // 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.
+// every package that the file belongs to, in every typechecking mode
+// applicable.
 func qualifiedObjsAtProtocolPos(ctx context.Context, s Snapshot, fh FileHandle, pp protocol.Position) ([]qualifiedObject, error) {
-	pkgs, err := s.PackagesForFile(ctx, fh.URI())
+	pkgs, err := s.PackagesForFile(ctx, fh.URI(), TypecheckAll)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/lsp/source/util.go b/internal/lsp/source/util.go
index 0d00849bb7a..d52f90b0b05 100644
--- a/internal/lsp/source/util.go
+++ b/internal/lsp/source/util.go
@@ -69,7 +69,7 @@ func (s mappedRange) URI() span.URI {
 // getParsedFile is a convenience function that extracts the Package and ParsedGoFile for a File in a Snapshot.
 // selectPackage is typically Narrowest/WidestPackageHandle below.
 func getParsedFile(ctx context.Context, snapshot Snapshot, fh FileHandle, selectPackage PackagePolicy) (Package, *ParsedGoFile, error) {
-	phs, err := snapshot.PackagesForFile(ctx, fh.URI())
+	phs, err := snapshot.PackagesForFile(ctx, fh.URI(), TypecheckWorkspace)
 	if err != nil {
 		return nil, nil, err
 	}
diff --git a/internal/lsp/source/view.go b/internal/lsp/source/view.go
index 8515220af62..e463491cb03 100644
--- a/internal/lsp/source/view.go
+++ b/internal/lsp/source/view.go
@@ -94,20 +94,21 @@ type Snapshot interface {
 	// BuiltinPackage returns information about the special builtin package.
 	BuiltinPackage(ctx context.Context) (*BuiltinPackage, error)
 
-	// PackagesForFile returns the packages that this file belongs to.
-	PackagesForFile(ctx context.Context, uri span.URI) ([]Package, error)
+	// PackagesForFile returns the packages that this file belongs to, checked
+	// in mode.
+	PackagesForFile(ctx context.Context, uri span.URI, mode TypecheckMode) ([]Package, error)
 
 	// GetActiveReverseDeps returns the active files belonging to the reverse
-	// dependencies of this file's package.
+	// dependencies of this file's package, checked in TypecheckWorkspace mode.
 	GetReverseDependencies(ctx context.Context, id string) ([]Package, error)
 
-	// CachedImportPaths returns all the imported packages loaded in this snapshot,
-	// indexed by their import path.
+	// CachedImportPaths returns all the imported packages loaded in this
+	// snapshot, indexed by their import path and checked in TypecheckWorkspace
+	// mode.
 	CachedImportPaths(ctx context.Context) (map[string]Package, error)
 
-	// KnownPackages returns all the packages loaded in this snapshot.
-	// Workspace packages may be parsed in ParseFull mode, whereas transitive
-	// dependencies will be in ParseExported mode.
+	// KnownPackages returns all the packages loaded in this snapshot, checked
+	// in TypecheckWorkspace mode.
 	KnownPackages(ctx context.Context) ([]Package, error)
 
 	// WorkspacePackages returns the snapshot's top-level packages.
@@ -323,7 +324,7 @@ type ParseMode int
 const (
 	// ParseHeader specifies that the main package declaration and imports are needed.
 	// This is the mode used when attempting to examine the package graph structure.
-	ParseHeader = ParseMode(iota)
+	ParseHeader ParseMode = iota
 
 	// ParseExported specifies that the public symbols are needed, but things like
 	// private symbols and function bodies are not.
@@ -337,6 +338,23 @@ const (
 	ParseFull
 )
 
+// TypecheckMode controls what kind of parsing should be done (see ParseMode)
+// while type checking a package.
+type TypecheckMode int
+
+const (
+	// Invalid default value.
+	TypecheckUnknown TypecheckMode = iota
+	// TypecheckFull means to use ParseFull.
+	TypecheckFull
+	// 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 {
 	FileHandle
 	Version() float64