internal/lsp: add support for document highlight

Change-Id: I232dbb0b66d690e45079808fd0dbf026c4459400
Reviewed-on: https://go-review.googlesource.com/c/tools/+/169277
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

go/packages/packagestest/expect.go

@@ -52,6 +52,7 @@ const (
 //   *regexp.Regexp : can only be supplied a regular expression literal
 //   token.Pos : has a file position calculated as described below.
 //   token.Position : has a file position calculated as described below.
+//   expect.Range: has a start and end position as described below.
 //   interface{} : will be passed any value
 //
 // Position calculation

internal/lsp/highlight.go

@@ -0,0 +1,24 @@
+// Copyright 2019 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 lsp
+
+import (
+	"golang.org/x/tools/internal/lsp/protocol"
+	"golang.org/x/tools/internal/span"
+)
+
+func toProtocolHighlight(m *protocol.ColumnMapper, spans []span.Span) []protocol.DocumentHighlight {
+	result := make([]protocol.DocumentHighlight, 0, len(spans))
+	kind := protocol.Text
+	for _, span := range spans {
+		r, err := m.Range(span)
+		if err != nil {
+			continue
+		}
+		h := protocol.DocumentHighlight{Kind: &kind, Range: r}
+		result = append(result, h)
+	}
+	return result
+}

internal/lsp/lsp_test.go

@@ -42,6 +42,7 @@ func testLSP(t *testing.T, exporter packagestest.Exporter) {
 	const expectedFormatCount = 4
 	const expectedDefinitionsCount = 16
 	const expectedTypeDefinitionsCount = 2
+	const expectedHighlightsCount = 2
 
 	files := packagestest.MustCopyFileTree(dir)
 	for fragment, operation := range files {
@@ -85,15 +86,17 @@ func testLSP(t *testing.T, exporter packagestest.Exporter) {
 	expectedFormat := make(formats)
 	expectedDefinitions := make(definitions)
 	expectedTypeDefinitions := make(definitions)
+	expectedHighlights := make(highlights)
 
 	// Collect any data that needs to be used by subsequent tests.
 	if err := exported.Expect(map[string]interface{}{
-		"diag":     expectedDiagnostics.collect,
-		"item":     completionItems.collect,
-		"complete": expectedCompletions.collect,
-		"format":   expectedFormat.collect,
-		"godef":    expectedDefinitions.collect,
-		"typdef":   expectedTypeDefinitions.collect,
+		"diag":      expectedDiagnostics.collect,
+		"item":      completionItems.collect,
+		"complete":  expectedCompletions.collect,
+		"format":    expectedFormat.collect,
+		"godef":     expectedDefinitions.collect,
+		"typdef":    expectedTypeDefinitions.collect,
+		"highlight": expectedHighlights.collect,
 	}); err != nil {
 		t.Fatal(err)
 	}
@@ -155,6 +158,16 @@ func testLSP(t *testing.T, exporter packagestest.Exporter) {
 		}
 		expectedTypeDefinitions.test(t, s, true)
 	})
+
+	t.Run("Highlights", func(t *testing.T) {
+		t.Helper()
+		if goVersion111 { // TODO(rstambler): Remove this when we no longer support Go 1.10.
+			if len(expectedHighlights) != expectedHighlightsCount {
+				t.Errorf("got %v highlights expected %v", len(expectedHighlights), expectedHighlightsCount)
+			}
+		}
+		expectedHighlights.test(t, s)
+	})
 }
 
 type diagnostics map[span.URI][]protocol.Diagnostic
@@ -162,6 +175,7 @@ type completionItems map[token.Pos]*protocol.CompletionItem
 type completions map[token.Position][]token.Pos
 type formats map[string]string
 type definitions map[protocol.Location]protocol.Location
+type highlights map[string][]protocol.Location
 
 func (d diagnostics) test(t *testing.T, v source.View) int {
 	count := 0
@@ -456,6 +470,39 @@ func (d definitions) collect(e *packagestest.Exported, fset *token.FileSet, src,
 	d[lSrc] = lTarget
 }
 
+func (h highlights) collect(e *packagestest.Exported, fset *token.FileSet, name string, rng packagestest.Range) {
+	s, m := testLocation(e, fset, rng)
+	loc, err := m.Location(s)
+	if err != nil {
+		return
+	}
+
+	h[name] = append(h[name], loc)
+}
+
+func (h highlights) test(t *testing.T, s *server) {
+	for name, locations := range h {
+		params := &protocol.TextDocumentPositionParams{
+			TextDocument: protocol.TextDocumentIdentifier{
+				URI: locations[0].URI,
+			},
+			Position: locations[0].Range.Start,
+		}
+		highlights, err := s.DocumentHighlight(context.Background(), params)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(highlights) != len(locations) {
+			t.Fatalf("got %d highlights for %s, expected %d", len(highlights), name, len(locations))
+		}
+		for i := range highlights {
+			if highlights[i].Range != locations[i].Range {
+				t.Errorf("want %v, got %v\n", locations[i].Range, highlights[i].Range)
+			}
+		}
+	}
+}
+
 func testLocation(e *packagestest.Exported, fset *token.FileSet, rng packagestest.Range) (span.Span, *protocol.ColumnMapper) {
 	spn, err := span.NewRange(fset, rng.Start, rng.End).Span()
 	if err != nil {

internal/lsp/server.go

@@ -127,6 +127,7 @@ func (s *server) Initialize(ctx context.Context, params *protocol.InitializePara
 				DocumentRangeFormattingProvider: true,
 				DocumentSymbolProvider:          true,
 				HoverProvider:                   true,
+				DocumentHighlightProvider:       true,
 				SignatureHelpProvider: &protocol.SignatureHelpOptions{
 					TriggerCharacters: []string{"(", ","},
 				},
@@ -423,8 +424,24 @@ func (s *server) References(context.Context, *protocol.ReferenceParams) ([]proto
 	return nil, notImplemented("References")
 }
 
-func (s *server) DocumentHighlight(context.Context, *protocol.TextDocumentPositionParams) ([]protocol.DocumentHighlight, error) {
-	return nil, notImplemented("DocumentHighlight")
+func (s *server) DocumentHighlight(ctx context.Context, params *protocol.TextDocumentPositionParams) ([]protocol.DocumentHighlight, error) {
+	f, m, err := newColumnMap(ctx, s.view, span.URI(params.TextDocument.URI))
+	if err != nil {
+		return nil, err
+	}
+
+	spn, err := m.PointSpan(params.Position)
+	if err != nil {
+		return nil, err
+	}
+
+	rng, err := spn.Range(m.Converter)
+	if err != nil {
+		return nil, err
+	}
+
+	spans := source.Highlight(ctx, f, rng.Start)
+	return toProtocolHighlight(m, spans), nil
 }
 
 func (s *server) DocumentSymbol(ctx context.Context, params *protocol.DocumentSymbolParams) ([]protocol.DocumentSymbol, error) {

internal/lsp/source/highlight.go

@@ -0,0 +1,42 @@
+// Copyright 2019 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 source
+
+import (
+	"context"
+	"go/ast"
+	"go/token"
+
+	"golang.org/x/tools/go/ast/astutil"
+	"golang.org/x/tools/internal/span"
+)
+
+func Highlight(ctx context.Context, f File, pos token.Pos) []span.Span {
+	fAST := f.GetAST(ctx)
+	fset := f.GetFileSet(ctx)
+	path, _ := astutil.PathEnclosingInterval(fAST, pos, pos)
+	if len(path) == 0 {
+		return nil
+	}
+
+	id, ok := path[0].(*ast.Ident)
+	if !ok {
+		return nil
+	}
+
+	var result []span.Span
+	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 {
+				s, err := nodeSpan(n, fset)
+				if err == nil {
+					result = append(result, s)
+				}
+			}
+			return true
+		})
+	}
+	return result
+}

internal/lsp/testdata/highlights/highlights.go

@@ -0,0 +1,15 @@
+package highlights
+
+import "fmt"
+
+type F struct{ bar int }
+
+var foo = F{bar: 52} //@highlight("foo", "foo")
+
+func Print() {
+	fmt.Println(foo) //@highlight("foo", "foo")
+}
+
+func (x *F) Inc() { //@highlight("x", "x")
+	x.bar++ //@highlight("x", "x")
+}