internal/lsp: improve highlighting functionality for fields and methods

Modified the way highlights are tested to allow for author to explicitly
mark the matches. Also added highlighting for fields and methods. Used
type checking in addition to ast to get better matching. Worked with
@stamblerre

Updates #34496

Change-Id: I462703e0011c4e0a4b98016e9c25af9bf1ead0b9
Reviewed-on: https://go-review.googlesource.com/c/tools/+/207899
Run-TryBot: Rohan Challa <rohan@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/internal/lsp/cmd/test/cmdtest.go b/internal/lsp/cmd/test/cmdtest.go
index 51b3d0a..3daf9b6 100644
--- a/internal/lsp/cmd/test/cmdtest.go
+++ b/internal/lsp/cmd/test/cmdtest.go
@@ -65,7 +65,7 @@
 	//TODO: add command line completions tests when it works
 }
 
-func (r *runner) Highlight(t *testing.T, name string, locations []span.Span) {
+func (r *runner) Highlight(t *testing.T, src span.Span, locations []span.Span) {
 	//TODO: add command line highlight tests when it works
 }
 
diff --git a/internal/lsp/lsp_test.go b/internal/lsp/lsp_test.go
index 1feddb1..a9a0d71 100644
--- a/internal/lsp/lsp_test.go
+++ b/internal/lsp/lsp_test.go
@@ -470,7 +470,7 @@
 	}
 }
 
-func (r *runner) Highlight(t *testing.T, name string, locations []span.Span) {
+func (r *runner) Highlight(t *testing.T, src span.Span, locations []span.Span) {
 	m, err := r.data.Mapper(locations[0].URI())
 	if err != nil {
 		t.Fatal(err)
@@ -491,7 +491,7 @@
 		t.Fatal(err)
 	}
 	if len(highlights) != len(locations) {
-		t.Fatalf("got %d highlights for %s, expected %d", len(highlights), name, len(locations))
+		t.Fatalf("got %d highlights for highlight at %v:%v:%v, expected %d", len(highlights), src.URI().Filename(), src.Start().Line(), src.Start().Column(), len(locations))
 	}
 	for i := range highlights {
 		if h, err := m.RangeSpan(highlights[i].Range); err != nil {
diff --git a/internal/lsp/source/highlight.go b/internal/lsp/source/highlight.go
index 8c0c104..d4fdc82 100644
--- a/internal/lsp/source/highlight.go
+++ b/internal/lsp/source/highlight.go
@@ -23,8 +23,27 @@
 	if err != nil {
 		return nil, err
 	}
-	fh := view.Snapshot().Handle(ctx, f)
-	ph := view.Session().Cache().ParseGoHandle(fh, ParseFull)
+	_, cphs, err := view.CheckPackageHandles(ctx, f)
+	if err != nil {
+		return nil, err
+	}
+	cph, err := WidestCheckPackageHandle(cphs)
+	if err != nil {
+		return nil, err
+	}
+	pkg, err := cph.Check(ctx)
+	if err != nil {
+		return nil, err
+	}
+	var ph ParseGoHandle
+	for _, file := range pkg.Files() {
+		if file.File().Identity().URI == f.URI() {
+			ph = file
+		}
+	}
+	if ph == nil {
+		return nil, errors.Errorf("no ParseGoHandle for %s", f.URI())
+	}
 	file, m, _, err := ph.Parse(ctx)
 	if err != nil {
 		return nil, err
@@ -41,22 +60,35 @@
 	if len(path) == 0 {
 		return nil, errors.Errorf("no enclosing position found for %v:%v", int(pos.Line), int(pos.Character))
 	}
+	result := []protocol.Range{}
 	id, ok := path[0].(*ast.Ident)
 	if !ok {
 		// If the cursor is not within an identifier, return empty results.
-		return []protocol.Range{}, nil
+		return result, nil
 	}
-	var result []protocol.Range
-	if id.Obj != nil {
-		ast.Inspect(path[len(path)-1], func(n ast.Node) bool {
-			if n, ok := n.(*ast.Ident); ok && n.Obj == id.Obj {
-				rng, err := nodeToProtocolRange(ctx, view, m, n)
-				if err == nil {
-					result = append(result, rng)
-				}
-			}
+	idObj := pkg.GetTypesInfo().ObjectOf(id)
+
+	ast.Inspect(path[len(path)-1], func(node ast.Node) bool {
+		n, ok := node.(*ast.Ident)
+		if !ok {
 			return true
-		})
-	}
+		}
+		if n.Name != id.Name {
+			return true
+		}
+		if n.Obj != id.Obj {
+			return true
+		}
+
+		nodeObj := pkg.GetTypesInfo().ObjectOf(n)
+		if nodeObj != idObj {
+			return false
+		}
+
+		if rng, err := nodeToProtocolRange(ctx, view, m, n); err == nil {
+			result = append(result, rng)
+		}
+		return true
+	})
 	return result, nil
 }
diff --git a/internal/lsp/source/source_test.go b/internal/lsp/source/source_test.go
index 8f84a26..d05ce8a 100644
--- a/internal/lsp/source/source_test.go
+++ b/internal/lsp/source/source_test.go
@@ -580,9 +580,8 @@
 	}
 }
 
-func (r *runner) Highlight(t *testing.T, name string, locations []span.Span) {
+func (r *runner) Highlight(t *testing.T, src span.Span, locations []span.Span) {
 	ctx := r.ctx
-	src := locations[0]
 	m, srcRng, err := spanToRange(r.data, src)
 	if err != nil {
 		t.Fatal(err)
@@ -592,7 +591,7 @@
 		t.Errorf("highlight failed for %s: %v", src.URI(), err)
 	}
 	if len(highlights) != len(locations) {
-		t.Errorf("got %d highlights for %s, expected %d", len(highlights), name, len(locations))
+		t.Errorf("got %d highlights for highlight at %v:%v:%v, expected %d", len(highlights), src.URI().Filename(), src.Start().Line(), src.Start().Column(), len(locations))
 	}
 	for i, got := range highlights {
 		want, err := m.Range(locations[i])
diff --git a/internal/lsp/testdata/highlights/highlights.go b/internal/lsp/testdata/highlights/highlights.go
index 9314842..dccd27e 100644
--- a/internal/lsp/testdata/highlights/highlights.go
+++ b/internal/lsp/testdata/highlights/highlights.go
@@ -1,15 +1,38 @@
 package highlights
 
-import "fmt"
+import (
+	"fmt"
+
+	"golang.org/x/tools/internal/lsp/protocol"
+)
 
 type F struct{ bar int }
 
-var foo = F{bar: 52} //@highlight("foo", "foo")
+var foo = F{bar: 52} //@mark(fooDeclaration, "foo"),highlight(fooDeclaration, fooDeclaration, fooUse)
 
-func Print() {
-	fmt.Println(foo) //@highlight("foo", "foo")
+func Print() { //@mark(printFunc, "Print"),highlight(printFunc, printFunc, printTest)
+	fmt.Println(foo) //@mark(fooUse, "foo"),highlight(fooUse, fooDeclaration, fooUse)
+	fmt.Print("yo")  //@mark(printSep, "Print"),highlight(printSep, printSep, print1, print2)
 }
 
-func (x *F) Inc() { //@highlight("x", "x")
-	x.bar++ //@highlight("x", "x")
+func (x *F) Inc() { //@mark(xDeclaration, "x"),highlight(xDeclaration, xDeclaration, xUse)
+	x.bar++ //@mark(xUse, "x"),highlight(xUse, xDeclaration, xUse)
+}
+
+func testFunctions() {
+	fmt.Print("main start") //@mark(print1, "Print"),highlight(print1, printSep, print1, print2)
+	fmt.Print("ok")         //@mark(print2, "Print"),highlight(print2, printSep, print1, print2)
+	Print()                 //@mark(printTest, "Print"),highlight(printTest, printFunc, printTest)
+}
+
+func toProtocolHighlight(rngs []protocol.Range) []protocol.DocumentHighlight { //@mark(doc1, "DocumentHighlight"),highlight(doc1, doc1, doc2, doc3)
+	result := make([]protocol.DocumentHighlight, 0, len(rngs)) //@mark(doc2, "DocumentHighlight"),highlight(doc2, doc1, doc2, doc3)
+	kind := protocol.Text
+	for _, rng := range rngs {
+		result = append(result, protocol.DocumentHighlight{ //@mark(doc3, "DocumentHighlight"),highlight(doc3, doc1, doc2, doc3)
+			Kind:  kind,
+			Range: rng,
+		})
+	}
+	return result
 }
diff --git a/internal/lsp/testdata/summary.txt.golden b/internal/lsp/testdata/summary.txt.golden
index 682605b..4410d51 100644
--- a/internal/lsp/testdata/summary.txt.golden
+++ b/internal/lsp/testdata/summary.txt.golden
@@ -13,7 +13,7 @@
 SuggestedFixCount = 1
 DefinitionsCount = 38
 TypeDefinitionsCount = 2
-HighlightsCount = 2
+HighlightsCount = 12
 ReferencesCount = 7
 RenamesCount = 22
 PrepareRenamesCount = 8
diff --git a/internal/lsp/tests/tests.go b/internal/lsp/tests/tests.go
index ba3aac0..2f1ea65 100644
--- a/internal/lsp/tests/tests.go
+++ b/internal/lsp/tests/tests.go
@@ -53,7 +53,7 @@
 type SuggestedFixes []span.Span
 type Definitions map[span.Span]Definition
 type Implementationses map[span.Span]Implementations
-type Highlights map[string][]span.Span
+type Highlights map[span.Span][]span.Span
 type References map[span.Span][]span.Span
 type Renames map[span.Span]string
 type PrepareRenames map[span.Span]*source.PrepareItem
@@ -113,7 +113,7 @@
 	SuggestedFix(*testing.T, span.Span)
 	Definition(*testing.T, span.Span, Definition)
 	Implementation(*testing.T, span.Span, Implementations)
-	Highlight(*testing.T, string, []span.Span)
+	Highlight(*testing.T, span.Span, []span.Span)
 	References(*testing.T, span.Span, []span.Span)
 	Rename(*testing.T, span.Span, string)
 	PrepareRename(*testing.T, span.Span, *source.PrepareItem)
@@ -479,10 +479,10 @@
 
 	t.Run("Highlight", func(t *testing.T) {
 		t.Helper()
-		for name, locations := range data.Highlights {
-			t.Run(name, func(t *testing.T) {
+		for pos, locations := range data.Highlights {
+			t.Run(spanName(pos), func(t *testing.T) {
 				t.Helper()
-				tests.Highlight(t, name, locations)
+				tests.Highlight(t, pos, locations)
 			})
 		}
 	})
@@ -826,8 +826,9 @@
 	data.Definitions[src] = d
 }
 
-func (data *Data) collectHighlights(name string, rng span.Span) {
-	data.Highlights[name] = append(data.Highlights[name], rng)
+func (data *Data) collectHighlights(src span.Span, expected []span.Span) {
+	// Declaring a highlight in a test file: @highlight(src, expected1, expected2)
+	data.Highlights[src] = append(data.Highlights[src], expected...)
 }
 
 func (data *Data) collectReferences(src span.Span, expected []span.Span) {