internal/lsp: fix find-references to search from all package variants

We previously searched the reverse dependencies of the "widest"
package that contained out starting identifier, but if our package has
tests then the widest package is the ".test" variant, and it has no
reverse dependencies. Fix by searching through all of the packages
that contain our starting identifier.

For example:

-- foo/foo.go --
package foo
func Foo() {}

-- foo/foo_test.go --
package foo
func TestFoo(t *testing.T) {}

-- bar/bar.go --
import "foo"
func _() {
  foo.Foo()
}

We would start searching from the foo.test variant, but we wouldn't
search package bar at all because bar does not import foo.test, it
imports plain foo. Now we search from both foo and foo.test (you still
need search foo.test to find references within foo_test.go).

Fixes golang/go#35936.

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

1 files changed, 43 insertions, 23 deletions

diff --git a/internal/lsp/references.go b/internal/lsp/references.go
index 675889bdf..0029264a0 100644
--- a/internal/lsp/references.go
+++ b/internal/lsp/references.go
@@ -29,48 +29,68 @@ func (s *Server) references(ctx context.Context, params *protocol.ReferenceParam
 	if f.Kind() != source.Go {
 		return nil, nil
 	}
-	ident, err := source.Identifier(ctx, snapshot, f, params.Position, source.WidestCheckPackageHandle)
+	phs, err := snapshot.PackageHandles(ctx, snapshot.Handle(ctx, f))
 	if err != nil {
 		return nil, nil
 	}
-	references, err := ident.References(ctx)
-	if err != nil {
-		log.Error(ctx, "no references", err, tag.Of("Identifier", ident.Name))
-	}
 
 	// Get the location of each reference to return as the result.
-	locations := make([]protocol.Location, 0, len(references))
-	seen := make(map[span.Span]bool)
-	for _, ref := range references {
-		refSpan, err := ref.Span()
+	var (
+		locations []protocol.Location
+		seen      = make(map[span.Span]bool)
+		lastIdent *source.IdentifierInfo
+	)
+	for _, ph := range phs {
+		ident, err := source.Identifier(ctx, snapshot, f, params.Position, source.SpecificPackageHandle(ph.ID()))
 		if err != nil {
-			return nil, err
+			if err == source.ErrNoIdentFound {
+				return nil, err
+			}
+			log.Error(ctx, "no identifier", err, tag.Of("Identifier", ident.Name))
+			continue
 		}
-		if seen[refSpan] {
-			continue // already added this location
-		}
-		seen[refSpan] = true
-		refRange, err := ref.Range()
+
+		lastIdent = ident
+
+		references, err := ident.References(ctx)
 		if err != nil {
-			return nil, err
+			log.Error(ctx, "no references", err, tag.Of("Identifier", ident.Name))
+			continue
+		}
+
+		for _, ref := range references {
+			refSpan, err := ref.Span()
+			if err != nil {
+				return nil, err
+			}
+			if seen[refSpan] {
+				continue // already added this location
+			}
+			seen[refSpan] = true
+			refRange, err := ref.Range()
+			if err != nil {
+				return nil, err
+			}
+			locations = append(locations, protocol.Location{
+				URI:   protocol.NewURI(ref.URI()),
+				Range: refRange,
+			})
 		}
-		locations = append(locations, protocol.Location{
-			URI:   protocol.NewURI(ref.URI()),
-			Range: refRange,
-		})
 	}
+
 	// Only add the identifier's declaration if the client requests it.
-	if params.Context.IncludeDeclaration {
-		rng, err := ident.Declaration.Range()
+	if params.Context.IncludeDeclaration && lastIdent != nil {
+		rng, err := lastIdent.Declaration.Range()
 		if err != nil {
 			return nil, err
 		}
 		locations = append([]protocol.Location{
 			{
-				URI:   protocol.NewURI(ident.Declaration.URI()),
+				URI:   protocol.NewURI(lastIdent.Declaration.URI()),
 				Range: rng,
 			},
 		}, locations...)
 	}
+
 	return locations, nil
 }