internal/lsp/source: use token.File-agnostic positions to dedupe refs

We were already using a token.File-agnostic position for object
positions in our references algorithm, but still relied on token.Pos for
identifying duplicate references. This breaks down when a file may have
multiple parsed representations in different packages.

While we should endeavor to eliminate duplicate parsing, our algorithms
should not rely on this for correctness. Update the reference
de-duplication to use the same position key as object search.

For golang/go#53796

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

internal/lsp/source/implementation.go

@@ -235,17 +235,23 @@ func qualifiedObjsAtProtocolPos(ctx context.Context, s Snapshot, uri span.URI, p
 	if err != nil {
 		return nil, err
 	}
-	return qualifiedObjsAtLocation(ctx, s, objSearchKey{uri, offset}, map[objSearchKey]bool{})
+	return qualifiedObjsAtLocation(ctx, s, positionKey{uri, offset}, map[positionKey]bool{})
 }
 
-type objSearchKey struct {
+// A positionKey identifies a byte offset within a file (URI).
+//
+// When a file has been parsed multiple times in the same FileSet,
+// there may be multiple token.Pos values denoting the same logical
+// position. In such situations, a positionKey may be used for
+// de-duplication.
+type positionKey struct {
 	uri    span.URI
 	offset int
 }
 
 // qualifiedObjsAtLocation finds all objects referenced at offset in uri, across
 // all packages in the snapshot.
-func qualifiedObjsAtLocation(ctx context.Context, s Snapshot, key objSearchKey, seen map[objSearchKey]bool) ([]qualifiedObject, error) {
+func qualifiedObjsAtLocation(ctx context.Context, s Snapshot, key positionKey, seen map[positionKey]bool) ([]qualifiedObject, error) {
 	if seen[key] {
 		return nil, nil
 	}
@@ -343,21 +349,8 @@ func qualifiedObjsAtLocation(ctx context.Context, s Snapshot, key objSearchKey,
 			// is in another package, but this should be good enough to find all
 			// uses.
 
-			pos := obj.Pos()
-			var uri span.URI
-			offset := -1
-			for _, pgf := range pkg.CompiledGoFiles() {
-				if pgf.Tok.Base() <= int(pos) && int(pos) <= pgf.Tok.Base()+pgf.Tok.Size() {
-					var err error
-					offset, err = safetoken.Offset(pgf.Tok, pos)
-					if err != nil {
-						return nil, err
-					}
-					uri = pgf.URI
-				}
-			}
-			if offset >= 0 {
-				otherObjs, err := qualifiedObjsAtLocation(ctx, s, objSearchKey{uri, offset}, seen)
+			if key, found := packagePositionKey(pkg, obj.Pos()); found {
+				otherObjs, err := qualifiedObjsAtLocation(ctx, s, key, seen)
 				if err != nil {
 					return nil, err
 				}
@@ -380,6 +373,19 @@ func qualifiedObjsAtLocation(ctx context.Context, s Snapshot, key objSearchKey,
 	return qualifiedObjs, nil
 }
 
+// packagePositionKey finds the positionKey for the given pos.
+//
+// The second result reports whether the position was found.
+func packagePositionKey(pkg Package, pos token.Pos) (positionKey, bool) {
+	for _, pgf := range pkg.CompiledGoFiles() {
+		offset, err := safetoken.Offset(pgf.Tok, pos)
+		if err == nil {
+			return positionKey{pgf.URI, offset}, true
+		}
+	}
+	return positionKey{}, false
+}
+
 // pathEnclosingObjNode returns the AST path to the object-defining
 // node associated with pos. "Object-defining" means either an
 // *ast.Ident mapped directly to a types.Object or an ast.Node mapped

internal/lsp/source/references.go

@@ -16,6 +16,7 @@ import (
 	"strconv"
 
 	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/lsp/bug"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/safetoken"
 	"golang.org/x/tools/internal/span"
@@ -127,7 +128,7 @@ func References(ctx context.Context, s Snapshot, f FileHandle, pp protocol.Posit
 func references(ctx context.Context, snapshot Snapshot, qos []qualifiedObject, includeDeclaration, includeInterfaceRefs, includeEmbeddedRefs bool) ([]*ReferenceInfo, error) {
 	var (
 		references []*ReferenceInfo
-		seen       = make(map[token.Pos]bool)
+		seen       = make(map[positionKey]bool)
 	)
 
 	pos := qos[0].obj.Pos()
@@ -189,10 +190,15 @@ func references(ctx context.Context, snapshot Snapshot, qos []qualifiedObject, i
 						continue
 					}
 				}
-				if seen[ident.Pos()] {
+				key, found := packagePositionKey(pkg, ident.Pos())
+				if !found {
+					bug.Reportf("ident %v (pos: %v) not found in package %v", ident.Name, ident.Pos(), pkg.Name())
 					continue
 				}
-				seen[ident.Pos()] = true
+				if seen[key] {
+					continue
+				}
+				seen[key] = true
 				rng, err := posToMappedRange(snapshot, pkg, ident.Pos(), ident.End())
 				if err != nil {
 					return nil, err