From df305b82d218535fccc89e230f23fef4c00302fe Mon Sep 17 00:00:00 2001
From: Rebecca Stambler <rstambler@golang.org>
Date: Wed, 4 Sep 2019 17:23:20 -0400
Subject: [PATCH] internal/lsp: fix declarations in references

Fixes golang/go#34087

Change-Id: I854c03cd124fe783f838dc53ee76cec5fffa563b
Reviewed-on: https://go-review.googlesource.com/c/tools/+/193379
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Ian Cottrell <iancottrell@google.com>
---
 internal/lsp/lsp_test.go                 |  1 -
 internal/lsp/references.go               | 21 +++++++++------
 internal/lsp/source/references.go        | 34 ++++++++++++------------
 internal/lsp/testdata/references/refs.go |  8 ++++++
 internal/lsp/tests/tests.go              |  2 +-
 5 files changed, 39 insertions(+), 27 deletions(-)

diff --git a/internal/lsp/lsp_test.go b/internal/lsp/lsp_test.go
index ba05f31578..ed7aa7cf09 100644
--- a/internal/lsp/lsp_test.go
+++ b/internal/lsp/lsp_test.go
@@ -621,7 +621,6 @@ func (r *runner) Reference(t *testing.T, data tests.References) {
 		if err != nil {
 			t.Fatalf("failed for %v: %v", src, err)
 		}
-
 		if len(got) != len(want) {
 			t.Errorf("references failed: different lengths got %v want %v", len(got), len(want))
 		}
diff --git a/internal/lsp/references.go b/internal/lsp/references.go
index 28d6587169..ea238d4a34 100644
--- a/internal/lsp/references.go
+++ b/internal/lsp/references.go
@@ -57,17 +57,22 @@ func (s *Server) references(ctx context.Context, params *protocol.ReferenceParam
 	// it is added to the beginning of the list if IncludeDeclaration
 	// was specified.
 	if params.Context.IncludeDeclaration {
-		rng, err := ident.Declaration.Range()
+		decSpan, err := ident.Declaration.Span()
 		if err != nil {
 			return nil, err
 		}
-		locations = append([]protocol.Location{
-			{
-				URI:   protocol.NewURI(ident.Declaration.URI()),
-				Range: rng,
-			},
-		}, locations...)
-
+		if !seen[decSpan] {
+			rng, err := ident.Declaration.Range()
+			if err != nil {
+				return nil, err
+			}
+			locations = append([]protocol.Location{
+				{
+					URI:   protocol.NewURI(ident.Declaration.URI()),
+					Range: rng,
+				},
+			}, locations...)
+		}
 	}
 	return locations, nil
 }
diff --git a/internal/lsp/source/references.go b/internal/lsp/source/references.go
index a81a36c850..bafd02ac4f 100644
--- a/internal/lsp/source/references.go
+++ b/internal/lsp/source/references.go
@@ -61,37 +61,37 @@ func (i *IdentifierInfo) References(ctx context.Context, view View) ([]*Referenc
 			if obj == nil || !sameObj(obj, i.Declaration.obj) {
 				continue
 			}
-			reference := &ReferenceInfo{
+			rng, err := posToRange(ctx, view, ident.Pos(), ident.End())
+			if err != nil {
+				return nil, err
+			}
+			// Add the declarations at the beginning of the references list.
+			references = append([]*ReferenceInfo{{
 				Name:          ident.Name,
 				ident:         ident,
 				obj:           obj,
 				pkg:           pkg,
 				isDeclaration: true,
-			}
-			if reference.mappedRange, err = posToRange(ctx, view, ident.Pos(), ident.End()); err != nil {
-				return nil, err
-			}
-			// Add the declarations at the beginning of the references list.
-			references = append([]*ReferenceInfo{reference}, references...)
+				mappedRange:   rng,
+			}}, references...)
 		}
 		for ident, obj := range info.Uses {
 			if obj == nil || !sameObj(obj, i.Declaration.obj) {
 				continue
 			}
-			reference := &ReferenceInfo{
-				Name:  ident.Name,
-				ident: ident,
-				pkg:   pkg,
-				obj:   obj,
-			}
-			if reference.mappedRange, err = posToRange(ctx, view, ident.Pos(), ident.End()); err != nil {
+			rng, err := posToRange(ctx, view, ident.Pos(), ident.End())
+			if err != nil {
 				return nil, err
 			}
-			references = append(references, reference)
+			references = append(references, &ReferenceInfo{
+				Name:        ident.Name,
+				ident:       ident,
+				pkg:         pkg,
+				obj:         obj,
+				mappedRange: rng,
+			})
 		}
-
 	}
-
 	return references, nil
 }
 
diff --git a/internal/lsp/testdata/references/refs.go b/internal/lsp/testdata/references/refs.go
index 89219c2bb6..a164715921 100644
--- a/internal/lsp/testdata/references/refs.go
+++ b/internal/lsp/testdata/references/refs.go
@@ -9,3 +9,11 @@ func _(_ int) []bool { //@mark(argInt, "int")
 func _(_ string) int { //@mark(returnInt, "int")
 	return 0
 }
+
+var q string //@mark(declQ, "q"),refs("q", declQ, assignQ, bobQ)
+
+func _() {
+	q = "hello" //@mark(assignQ, "q")
+	bob := func(_ string) {}
+	bob(q) //@mark(bobQ, "q")
+}
diff --git a/internal/lsp/tests/tests.go b/internal/lsp/tests/tests.go
index ac2e0c268c..956a9226e8 100644
--- a/internal/lsp/tests/tests.go
+++ b/internal/lsp/tests/tests.go
@@ -38,7 +38,7 @@ const (
 	ExpectedTypeDefinitionsCount   = 2
 	ExpectedFoldingRangesCount     = 1
 	ExpectedHighlightsCount        = 2
-	ExpectedReferencesCount        = 5
+	ExpectedReferencesCount        = 6
 	ExpectedRenamesCount           = 20
 	ExpectedPrepareRenamesCount    = 8
 	ExpectedSymbolsCount           = 1