internal/lsp/source: fix Highlight for std and 3rd-party packages

Don't use GetParsedFile because it extracts a package with
TypecheckWorkspace mode, but we want to support std and
3rd-party packages. So, we reuse the content of GetParsedFile
directly in Highlight function with TypecheckFull mode.

Fixes golang/go#43511

Change-Id: Ibd1d42e28a6739ba011113df0e6e7e98d1b86eb5
GitHub-Last-Rev: 3746092
GitHub-Pull-Request: golang#298
Reviewed-on: https://go-review.googlesource.com/c/tools/+/307171
Trust: Robert Findley <rfindley@google.com>
Trust: Rebecca Stambler <rstambler@golang.org>
Reviewed-by: 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>

gopls/internal/regtest/misc/highlight_test.go

@@ -0,0 +1,151 @@
+// Copyright 2021 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package misc
+
+import (
+	"sort"
+	"testing"
+
+	. "golang.org/x/tools/gopls/internal/regtest"
+	"golang.org/x/tools/internal/lsp/fake"
+	"golang.org/x/tools/internal/lsp/protocol"
+)
+
+func TestWorkspacePackageHighlight(t *testing.T) {
+	const mod = `
+-- go.mod --
+module mod.com
+
+go 1.12
+-- main.go --
+package main
+
+func main() {
+	var A string = "A"
+	x := "x-" + A
+	println(A, x)
+}`
+
+	Run(t, mod, func(t *testing.T, env *Env) {
+		const file = "main.go"
+		env.OpenFile(file)
+		_, pos := env.GoToDefinition(file, env.RegexpSearch(file, `var (A) string`))
+
+		checkHighlights(env, file, pos, 3)
+	})
+}
+
+func TestStdPackageHighlight_Issue43511(t *testing.T) {
+	const mod = `
+-- go.mod --
+module mod.com
+
+go 1.12
+-- main.go --
+package main
+
+import "fmt"
+
+func main() {
+	fmt.Printf()
+}`
+
+	Run(t, mod, func(t *testing.T, env *Env) {
+		env.OpenFile("main.go")
+		file, _ := env.GoToDefinition("main.go", env.RegexpSearch("main.go", `fmt\.(Printf)`))
+		pos := env.RegexpSearch(file, `func Printf\((format) string`)
+
+		checkHighlights(env, file, pos, 2)
+	})
+}
+
+func TestThirdPartyPackageHighlight_Issue43511(t *testing.T) {
+	const proxy = `
+-- example.com@v1.2.3/go.mod --
+module example.com
+
+go 1.12
+-- example.com@v1.2.3/global/global.go --
+package global
+
+const A = 1
+
+func foo() {
+	_ = A
+}
+
+func bar() int {
+	return A + A
+}
+-- example.com@v1.2.3/local/local.go --
+package local
+
+func foo() int {
+	const b = 2
+
+	return b * b * (b+1) + b
+}`
+
+	const mod = `
+-- go.mod --
+module mod.com
+
+go 1.12
+
+require example.com v1.2.3
+-- go.sum --
+example.com v1.2.3 h1:WFzrgiQJwEDJNLDUOV1f9qlasQkvzXf2UNLaNIqbWsI=
+example.com v1.2.3/go.mod h1:Y2Rc5rVWjWur0h3pd9aEvK5Pof8YKDANh9gHA2Maujo=
+-- main.go --
+package main
+
+import (
+	_ "example.com/global"
+	_ "example.com/local"
+)
+
+func main() {}`
+
+	WithOptions(
+		ProxyFiles(proxy),
+	).Run(t, mod, func(t *testing.T, env *Env) {
+		env.OpenFile("main.go")
+
+		file, _ := env.GoToDefinition("main.go", env.RegexpSearch("main.go", `"example.com/global"`))
+		pos := env.RegexpSearch(file, `const (A)`)
+		checkHighlights(env, file, pos, 4)
+
+		file, _ = env.GoToDefinition("main.go", env.RegexpSearch("main.go", `"example.com/local"`))
+		pos = env.RegexpSearch(file, `const (b)`)
+		checkHighlights(env, file, pos, 5)
+	})
+}
+
+func checkHighlights(env *Env, file string, pos fake.Pos, highlightCount int) {
+	t := env.T
+	t.Helper()
+
+	highlights := env.DocumentHighlight(file, pos)
+	if len(highlights) != highlightCount {
+		t.Fatalf("expected %v highlight(s), got %v", highlightCount, len(highlights))
+	}
+
+	references := env.References(file, pos)
+	if len(highlights) != len(references) {
+		t.Fatalf("number of highlights and references is expected to be equal: %v != %v", len(highlights), len(references))
+	}
+
+	sort.Slice(highlights, func(i, j int) bool {
+		return protocol.CompareRange(highlights[i].Range, highlights[j].Range) < 0
+	})
+	sort.Slice(references, func(i, j int) bool {
+		return protocol.CompareRange(references[i].Range, references[j].Range) < 0
+	})
+	for i := range highlights {
+		if highlights[i].Range != references[i].Range {
+			t.Errorf("highlight and reference ranges are expected to be equal: %v != %v", highlights[i].Range, references[i].Range)
+		}
+	}
+}

gopls/internal/regtest/wrappers.go

@@ -230,6 +230,15 @@ func (e *Env) DocumentLink(name string) []protocol.DocumentLink {
 	return links
 }
 
+func (e *Env) DocumentHighlight(name string, pos fake.Pos) []protocol.DocumentHighlight {
+	e.T.Helper()
+	highlights, err := e.Editor.DocumentHighlight(e.Ctx, name, pos)
+	if err != nil {
+		e.T.Fatal(err)
+	}
+	return highlights
+}
+
 func checkIsFatal(t *testing.T, err error) {
 	t.Helper()
 	if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrClosedPipe) {

internal/lsp/fake/editor.go

@@ -1094,3 +1094,17 @@ func (e *Editor) DocumentLink(ctx context.Context, path string) ([]protocol.Docu
 	params.TextDocument.URI = e.sandbox.Workdir.URI(path)
 	return e.Server.DocumentLink(ctx, params)
 }
+
+func (e *Editor) DocumentHighlight(ctx context.Context, path string, pos Pos) ([]protocol.DocumentHighlight, error) {
+	if e.Server == nil {
+		return nil, nil
+	}
+	if err := e.checkBufferPosition(path, pos); err != nil {
+		return nil, err
+	}
+	params := &protocol.DocumentHighlightParams{}
+	params.TextDocument.URI = e.sandbox.Workdir.URI(path)
+	params.Position = pos.ToProtocolPosition()
+
+	return e.Server.DocumentHighlight(ctx, params)
+}

internal/lsp/source/highlight.go

@@ -22,10 +22,18 @@ func Highlight(ctx context.Context, snapshot Snapshot, fh FileHandle, pos protoc
 	ctx, done := event.Start(ctx, "source.Highlight")
 	defer done()
 
-	pkg, pgf, err := GetParsedFile(ctx, snapshot, fh, WidestPackage)
+	// Don't use GetParsedFile because it uses TypecheckWorkspace, and we
+	// always want fully parsed files for highlight, regardless of whether
+	// the file belongs to a workspace package.
+	pkg, err := snapshot.PackageForFile(ctx, fh.URI(), TypecheckFull, WidestPackage)
+	if err != nil {
+		return nil, errors.Errorf("getting package for Highlight: %w", err)
+	}
+	pgf, err := pkg.File(fh.URI())
 	if err != nil {
 		return nil, errors.Errorf("getting file for Highlight: %w", err)
 	}
+
 	spn, err := pgf.Mapper.PointSpan(pos)
 	if err != nil {
 		return nil, err