gopls/internal/lsp/source: eliminate a couple uses of posToMappedRange

Eliminate a couple uses of posToMappedRange, which potentially
type-checks, where it is clearly unnecessary.

Also improve test output for highlight.

Updates golang/go#57987
Updates golang/go#54845

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

gopls/internal/lsp/lsp_test.go

@@ -771,45 +771,56 @@ func (r *runner) Implementation(t *testing.T, spn span.Span, wantSpans []span.Sp
 	}
 }
 
-func (r *runner) Highlight(t *testing.T, src span.Span, locations []span.Span) {
+func (r *runner) Highlight(t *testing.T, src span.Span, spans []span.Span) {
 	m, err := r.data.Mapper(src.URI())
 	if err != nil {
 		t.Fatal(err)
 	}
 	loc, err := m.SpanLocation(src)
 	if err != nil {
-		t.Fatalf("failed for %v: %v", locations[0], err)
+		t.Fatal(err)
 	}
 	tdpp := protocol.TextDocumentPositionParams{
-		TextDocument: protocol.TextDocumentIdentifier{URI: loc.URI},
+		TextDocument: protocol.TextDocumentIdentifier{
+			URI: loc.URI,
+		},
 		Position: loc.Range.Start,
 	}
+	if err != nil {
+		t.Fatalf("Mapper.SpanDocumentPosition(%v) failed: %v", src, err)
+	}
 	params := &protocol.DocumentHighlightParams{
 		TextDocumentPositionParams: tdpp,
 	}
 	highlights, err := r.server.DocumentHighlight(r.ctx, params)
 	if err != nil {
-		t.Fatal(err)
+		t.Fatalf("DocumentHighlight(%v) failed: %v", params, err)
 	}
-	if len(highlights) != 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))
+	var got []protocol.Range
+	for _, h := range highlights {
+		got = append(got, h.Range)
 	}
-	// Check to make sure highlights have a valid range.
-	var results []span.Span
-	for i := range highlights {
-		h, err := m.RangeSpan(highlights[i].Range)
+
+	var want []protocol.Range
+	for _, s := range spans {
+		rng, err := m.SpanRange(s)
 		if err != nil {
-			t.Fatalf("failed for %v: %v", highlights[i], err)
+			t.Fatalf("Mapper.SpanRange(%v) failed: %v", s, err)
 		}
-		results = append(results, h)
+		want = append(want, rng)
 	}
-	// Sort results to make tests deterministic since DocumentHighlight uses a map.
-	span.SortSpans(results)
-	// Check to make sure all the expected highlights are found.
-	for i := range results {
-		if results[i] != locations[i] {
-			t.Errorf("want %v, got %v\n", locations[i], results[i])
+
+	sortRanges := func(s []protocol.Range) {
+		sort.Slice(s, func(i, j int) bool {
+			return protocol.CompareRange(s[i], s[j]) < 0
+		})
 	}
+
+	sortRanges(got)
+	sortRanges(want)
+
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Errorf("DocumentHighlight(%v) mismatch (-want +got):\n%s", src, diff)
 	}
 }
 

gopls/internal/lsp/source/highlight.go

@@ -48,28 +48,28 @@ func Highlight(ctx context.Context, snapshot Snapshot, fh FileHandle, position p
 			}
 		}
 	}
-	result, err := highlightPath(pkg, path)
+	result, err := highlightPath(path, pgf.File, pkg.GetTypesInfo())
 	if err != nil {
 		return nil, err
 	}
 	var ranges []protocol.Range
 	for rng := range result {
-		mRng, err := posToMappedRange(ctx, snapshot, pkg, rng.start, rng.end)
+		rng, err := pgf.PosRange(rng.start, rng.end)
 		if err != nil {
 			return nil, err
 		}
-		ranges = append(ranges, mRng.Range())
+		ranges = append(ranges, rng)
 	}
 	return ranges, nil
 }
 
-func highlightPath(pkg Package, path []ast.Node) (map[posRange]struct{}, error) {
+func highlightPath(path []ast.Node, file *ast.File, info *types.Info) (map[posRange]struct{}, error) {
 	result := make(map[posRange]struct{})
 	switch node := path[0].(type) {
 	case *ast.BasicLit:
 		if len(path) > 1 {
 			if _, ok := path[1].(*ast.ImportSpec); ok {
-				err := highlightImportUses(pkg, path, result)
+				err := highlightImportUses(path, info, result)
 				return result, err
 			}
 		}
@@ -77,7 +77,9 @@ func highlightPath(pkg Package, path []ast.Node) (map[posRange]struct{}, error)
 	case *ast.ReturnStmt, *ast.FuncDecl, *ast.FuncType:
 		highlightFuncControlFlow(path, result)
 	case *ast.Ident:
-		highlightIdentifiers(pkg, path, result)
+		// Check if ident is inside return or func decl.
+		highlightFuncControlFlow(path, result)
+		highlightIdentifier(node, file, info, result)
 	case *ast.ForStmt, *ast.RangeStmt:
 		highlightLoopControlFlow(path, result)
 	case *ast.SwitchStmt:
@@ -426,7 +428,7 @@ func labelDecl(n *ast.Ident) *ast.Ident {
 	return stmt.Label
 }
 
-func highlightImportUses(pkg Package, path []ast.Node, result map[posRange]struct{}) error {
+func highlightImportUses(path []ast.Node, info *types.Info, result map[posRange]struct{}) error {
 	basicLit, ok := path[0].(*ast.BasicLit)
 	if !ok {
 		return fmt.Errorf("highlightImportUses called with an ast.Node of type %T", basicLit)
@@ -440,7 +442,7 @@ func highlightImportUses(pkg Package, path []ast.Node, result map[posRange]struc
 		if !ok {
 			return true
 		}
-		obj, ok := pkg.GetTypesInfo().ObjectOf(n).(*types.PkgName)
+		obj, ok := info.ObjectOf(n).(*types.PkgName)
 		if !ok {
 			return true
 		}
@@ -453,19 +455,16 @@ func highlightImportUses(pkg Package, path []ast.Node, result map[posRange]struc
 	return nil
 }
 
-func highlightIdentifiers(pkg Package, path []ast.Node, result map[posRange]struct{}) error {
-	id, ok := path[0].(*ast.Ident)
-	if !ok {
-		return fmt.Errorf("highlightIdentifiers called with an ast.Node of type %T", id)
-	}
-	// Check if ident is inside return or func decl.
-	highlightFuncControlFlow(path, result)
-
-	// TODO: maybe check if ident is a reserved word, if true then don't continue and return results.
-
-	idObj := pkg.GetTypesInfo().ObjectOf(id)
+func highlightIdentifier(id *ast.Ident, file *ast.File, info *types.Info, result map[posRange]struct{}) {
+	// TODO(rfindley): idObj may be nil. Note that returning early in this case
+	// causes tests to fail (because the nObj == idObj check below was succeeded
+	// for nil == nil!)
+	//
+	// Revisit this. If ObjectOf is nil, there are type errors, and it seems
+	// reasonable for identifier highlighting not to work.
+	idObj := info.ObjectOf(id)
 	pkgObj, isImported := idObj.(*types.PkgName)
-	ast.Inspect(path[len(path)-1], func(node ast.Node) bool {
+	ast.Inspect(file, func(node ast.Node) bool {
 		if imp, ok := node.(*ast.ImportSpec); ok && isImported {
 			highlightImport(pkgObj, imp, result)
 		}
@@ -476,12 +475,11 @@ func highlightIdentifiers(pkg Package, path []ast.Node, result map[posRange]stru
 		if n.Name != id.Name {
 			return false
 		}
-		if nObj := pkg.GetTypesInfo().ObjectOf(n); nObj == idObj {
+		if nObj := info.ObjectOf(n); nObj == idObj {
 			result[posRange{start: n.Pos(), end: n.End()}] = struct{}{}
 		}
 		return false
 	})
-	return nil
 }
 
 func highlightImport(obj *types.PkgName, imp *ast.ImportSpec, result map[posRange]struct{}) {

gopls/internal/lsp/source/references.go

@@ -141,7 +141,12 @@ func references(ctx context.Context, snapshot Snapshot, qos []qualifiedObject, i
 					continue
 				}
 				seen[key] = true
-				rng, err := posToMappedRange(ctx, snapshot, pkg, ident.Pos(), ident.End())
+				filename := pkg.FileSet().File(ident.Pos()).Name()
+				pgf, err := pkg.File(span.URIFromPath(filename))
+				if err != nil {
+					return nil, err
+				}
+				rng, err := pgf.PosMappedRange(ident.Pos(), ident.End())
 				if err != nil {
 					return nil, err
 				}