From 1f763dfd2e10ac7b454ff87f33c2b31690c9ab3e Mon Sep 17 00:00:00 2001
From: pjw <pjw@google.com>
Date: Fri, 1 Apr 2022 11:10:20 -0400
Subject: [PATCH] internal/lsp: add semantic tokens for arrows in declarations

Fixes: golang.org/#52024

Change-Id: Ib79a839f9bdeaef1250925add69dea3ae32e1cae
Reviewed-on: https://go-review.googlesource.com/c/tools/+/397479
Reviewed-by: Hyang-Ah Hana Kim <hyangah@gmail.com>
Trust: Hyang-Ah Hana Kim <hyangah@gmail.com>
Trust: Peter Weinberger <pjw@google.com>
---
 internal/lsp/cmd/test/semanticdriver.go    |  4 +++-
 internal/lsp/semantic.go                   | 13 +++++++++----
 internal/lsp/testdata/semantic/a.go.golden |  4 ++--
 internal/lsp/testdata/semantic/b.go        |  4 ++++
 internal/lsp/testdata/semantic/b.go.golden |  4 ++++
 5 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/internal/lsp/cmd/test/semanticdriver.go b/internal/lsp/cmd/test/semanticdriver.go
index 80dc61e3d6..247f755bf2 100644
--- a/internal/lsp/cmd/test/semanticdriver.go
+++ b/internal/lsp/cmd/test/semanticdriver.go
@@ -27,7 +27,9 @@ func (r *runner) SemanticTokens(t *testing.T, spn span.Span) {
 		t.Errorf("want(%d-%d) != got(%d-%d) for %s", len(want), len(lwant), len(got), len(lgot), r.Normalize(filename))
 		for i := 0; i < len(lwant) && i < len(lgot); i++ {
 			if lwant[i] != lgot[i] {
-				t.Errorf("line %d:\nwant%q\ngot %q\n", i, lwant[i], lgot[i])
+				// This is the line number in the golden file.
+				// It is one larger than the line number in the source file.
+				t.Errorf("line %d:\nwant%q\ngot %q\n", i+2, lwant[i], lgot[i])
 			}
 		}
 	}
diff --git a/internal/lsp/semantic.go b/internal/lsp/semantic.go
index 4318b7baba..4c9e5d31a7 100644
--- a/internal/lsp/semantic.go
+++ b/internal/lsp/semantic.go
@@ -328,12 +328,17 @@ func (e *encoded) inspector(n ast.Node) bool {
 		e.token(x.Case, len(iam), tokKeyword, nil)
 	case *ast.ChanType:
 		// chan | chan <- | <- chan
-		if x.Arrow == token.NoPos || x.Arrow != x.Begin {
+		switch {
+		case x.Arrow == token.NoPos:
 			e.token(x.Begin, len("chan"), tokKeyword, nil)
-			break
+		case x.Arrow == x.Begin:
+			e.token(x.Arrow, 2, tokOperator, nil)
+			pos := e.findKeyword("chan", x.Begin+2, x.Value.Pos())
+			e.token(pos, len("chan"), tokKeyword, nil)
+		case x.Arrow != x.Begin:
+			e.token(x.Begin, len("chan"), tokKeyword, nil)
+			e.token(x.Arrow, 2, tokOperator, nil)
 		}
-		pos := e.findKeyword("chan", x.Begin+2, x.Value.Pos())
-		e.token(pos, len("chan"), tokKeyword, nil)
 	case *ast.CommClause:
 		iam := len("case")
 		if x.Comm == nil {
diff --git a/internal/lsp/testdata/semantic/a.go.golden b/internal/lsp/testdata/semantic/a.go.golden
index 4622ae4d74..19dd412407 100644
--- a/internal/lsp/testdata/semantic/a.go.golden
+++ b/internal/lsp/testdata/semantic/a.go.golden
@@ -13,8 +13,8 @@
 	/*⇒1,variable,[definition]*/a           = /*⇒3,namespace,[]*/fmt./*⇒5,function,[]*/Print
 	/*⇒1,variable,[definition]*/b  []/*⇒6,type,[defaultLibrary]*/string = []/*⇒6,type,[defaultLibrary]*/string{/*⇒5,string,[]*/"foo"}
 	/*⇒2,variable,[definition]*/c1 /*⇒4,keyword,[]*/chan /*⇒3,type,[defaultLibrary]*/int
-	/*⇒2,variable,[definition]*/c2 <-/*⇒4,keyword,[]*/chan /*⇒3,type,[defaultLibrary]*/int
-	/*⇒2,variable,[definition]*/c3 = /*⇒4,function,[defaultLibrary]*/make([]/*⇒4,keyword,[]*/chan<- /*⇒3,type,[defaultLibrary]*/int)
+	/*⇒2,variable,[definition]*/c2 /*⇒2,operator,[]*/<-/*⇒4,keyword,[]*/chan /*⇒3,type,[defaultLibrary]*/int
+	/*⇒2,variable,[definition]*/c3 = /*⇒4,function,[defaultLibrary]*/make([]/*⇒4,keyword,[]*/chan/*⇒2,operator,[]*/<- /*⇒3,type,[defaultLibrary]*/int)
 	/*⇒1,variable,[definition]*/b  = /*⇒1,type,[]*/A{/*⇒1,variable,[]*/X: /*⇒2,number,[]*/23}
 	/*⇒1,variable,[definition]*/m  /*⇒3,keyword,[]*/map[/*⇒4,type,[defaultLibrary]*/bool][/*⇒1,number,[]*/3]/*⇒1,operator,[]*/*/*⇒7,type,[defaultLibrary]*/float64
 )
diff --git a/internal/lsp/testdata/semantic/b.go b/internal/lsp/testdata/semantic/b.go
index a536f97bfe..496b0863d7 100644
--- a/internal/lsp/testdata/semantic/b.go
+++ b/internal/lsp/testdata/semantic/b.go
@@ -32,3 +32,7 @@ type CC struct {
 }
 type D func(aa AA) (BB error)
 type E func(AA) BB
+
+var a chan<- chan int
+var b chan<- <-chan int
+var c <-chan <-chan int
diff --git a/internal/lsp/testdata/semantic/b.go.golden b/internal/lsp/testdata/semantic/b.go.golden
index 203f6b1893..59071374ba 100644
--- a/internal/lsp/testdata/semantic/b.go.golden
+++ b/internal/lsp/testdata/semantic/b.go.golden
@@ -34,3 +34,7 @@
 /*⇒4,keyword,[]*/type /*⇒1,type,[definition]*/D /*⇒4,keyword,[]*/func(/*⇒2,parameter,[definition]*/aa /*⇒2,type,[]*/AA) (/*⇒2,parameter,[definition]*/BB /*⇒5,type,[]*/error)
 /*⇒4,keyword,[]*/type /*⇒1,type,[definition]*/E /*⇒4,keyword,[]*/func(/*⇒2,type,[]*/AA) /*⇒2,type,[]*/BB
 
+/*⇒3,keyword,[]*/var /*⇒1,variable,[definition]*/a /*⇒4,keyword,[]*/chan/*⇒2,operator,[]*/<- /*⇒4,keyword,[]*/chan /*⇒3,type,[defaultLibrary]*/int
+/*⇒3,keyword,[]*/var /*⇒1,variable,[definition]*/b /*⇒4,keyword,[]*/chan/*⇒2,operator,[]*/<- /*⇒2,operator,[]*/<-/*⇒4,keyword,[]*/chan /*⇒3,type,[defaultLibrary]*/int
+/*⇒3,keyword,[]*/var /*⇒1,variable,[definition]*/c /*⇒2,operator,[]*/<-/*⇒4,keyword,[]*/chan /*⇒2,operator,[]*/<-/*⇒4,keyword,[]*/chan /*⇒3,type,[defaultLibrary]*/int
+