internal/lsp: fix control flow highlighting taking precedence

The control flow highlighting is taking precedence when
you are highlighting a key:value expression within the return statement.
Expected behavior is to just highlight all instances of the key or value and ignore
the control flow statement when inside the scope.

Fixes #36057

Change-Id: If4b254151c38d152f337833c55a456f8dce18be7
Reviewed-on: https://go-review.googlesource.com/c/tools/+/210558
Reviewed-by: Rebecca Stambler <[email protected]>
Run-TryBot: Rohan Challa <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>

Author: ridersofrohan

internal/lsp/source/highlight.go

@@ -62,6 +62,9 @@ Outer:
 	// Reverse walk the path till we get to the func block.
 	for _, n := range path {
 		switch node := n.(type) {
+		case *ast.KeyValueExpr:
+			// If cursor is in a key: value expr, we don't want control flow highlighting
+			return nil, nil
 		case *ast.Field:
 			inReturnList = true
 		case *ast.FuncLit:
@@ -232,7 +235,7 @@ func highlightIdentifiers(ctx context.Context, snapshot Snapshot, m *protocol.Co
 		if !ok {
 			return true
 		}
-		if n.Name != id.Name || n.Obj != id.Obj {
+		if n.Name != id.Name {
 			return false
 		}
 		if nObj := pkg.GetTypesInfo().ObjectOf(n); nObj != idObj {

internal/lsp/testdata/highlights/highlights.go

@@ -7,17 +7,23 @@ import (
 	"golang.org/x/tools/internal/lsp/protocol"
 )
 
-type F struct{ bar int }
+type F struct{ bar int } //@mark(barDeclaration, "bar"),highlight(barDeclaration, barDeclaration, bar1, bar2, bar3)
 
-var foo = F{bar: 52} //@mark(fooDeclaration, "foo"),highlight(fooDeclaration, fooDeclaration, fooUse)
+func _() F {
+	return F{
+		bar: 123, //@mark(bar1, "bar"),highlight(bar1, barDeclaration, bar1, bar2, bar3)
+	}
+}
+
+var foo = F{bar: 52} //@mark(fooDeclaration, "foo"),mark(bar2, "bar"),highlight(fooDeclaration, fooDeclaration, fooUse),highlight(bar2, barDeclaration, bar1, bar2, bar3)
 
 func Print() { //@mark(printFunc, "Print"),highlight(printFunc, printFunc, printTest)
 	fmt.Println(foo) //@mark(fooUse, "foo"),highlight(fooUse, fooDeclaration, fooUse)
 	fmt.Print("yo")  //@mark(printSep, "Print"),highlight(printSep, printSep, print1, print2)
 }
 
 func (x *F) Inc() { //@mark(xDeclaration, "x"),highlight(xDeclaration, xDeclaration, xUse)
-	x.bar++ //@mark(xUse, "x"),highlight(xUse, xDeclaration, xUse)
+	x.bar++ //@mark(xUse, "x"),mark(bar3, "bar"),highlight(xUse, xDeclaration, xUse),highlight(bar3, barDeclaration, bar1, bar2, bar3)
 }
 
 func testFunctions() {

internal/lsp/testdata/summary.txt.golden

@@ -13,7 +13,7 @@ ImportCount = 7
 SuggestedFixCount = 1
 DefinitionsCount = 38
 TypeDefinitionsCount = 2
-HighlightsCount = 37
+HighlightsCount = 41
 ReferencesCount = 7
 RenamesCount = 22
 PrepareRenamesCount = 8