internal/lsp: add inlay hints for variable types

For golang/go#52343.
For golang/vscode-go#1631.

Change-Id: I94a1b3c389d8bfaa48754e28a52ef76c29eb6ead
Reviewed-on: https://go-review.googlesource.com/c/tools/+/411100
Run-TryBot: Jamal Carvalho <jamal@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Hyang-Ah Hana Kim <hyangah@gmail.com>
Reviewed-by: Suzy Mueller <suzmue@golang.org>

internal/lsp/source/inlay_hint.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"fmt"
 	"go/ast"
+	"go/token"
 	"go/types"
 
 	"golang.org/x/tools/internal/event"
@@ -30,12 +31,17 @@ func InlayHint(ctx context.Context, snapshot Snapshot, fh FileHandle, _ protocol
 
 	tmap := lsppos.NewTokenMapper(pgf.Src, pgf.Tok)
 	info := pkg.GetTypesInfo()
+	q := Qualifier(pgf.File, pkg.GetTypes(), info)
 
 	var hints []protocol.InlayHint
 	ast.Inspect(pgf.File, func(node ast.Node) bool {
 		switch n := node.(type) {
 		case *ast.CallExpr:
 			hints = append(hints, parameterNames(n, tmap, info)...)
+		case *ast.AssignStmt:
+			hints = append(hints, assignVariableTypes(n, tmap, info, &q)...)
+		case *ast.RangeStmt:
+			hints = append(hints, rangeVariableTypes(n, tmap, info, &q)...)
 		}
 		return true
 	})
@@ -78,6 +84,47 @@ func parameterNames(node *ast.CallExpr, tmap *lsppos.TokenMapper, info *types.In
 	return hints
 }
 
+func assignVariableTypes(node *ast.AssignStmt, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) []protocol.InlayHint {
+	if node.Tok != token.DEFINE {
+		return nil
+	}
+	var hints []protocol.InlayHint
+	for _, v := range node.Lhs {
+		if h := variableType(v, tmap, info, q); h != nil {
+			hints = append(hints, *h)
+		}
+	}
+	return hints
+}
+
+func rangeVariableTypes(node *ast.RangeStmt, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) []protocol.InlayHint {
+	var hints []protocol.InlayHint
+	if h := variableType(node.Key, tmap, info, q); h != nil {
+		hints = append(hints, *h)
+	}
+	if h := variableType(node.Value, tmap, info, q); h != nil {
+		hints = append(hints, *h)
+	}
+	return hints
+}
+
+func variableType(e ast.Expr, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) *protocol.InlayHint {
+	typ := info.TypeOf(e)
+	if typ == nil {
+		return nil
+	}
+	end, ok := tmap.Position(e.End())
+	if !ok {
+		return nil
+	}
+	return &protocol.InlayHint{
+		Position:    &end,
+		Label:       buildLabel(types.TypeString(typ, *q)),
+		Kind:        protocol.Type,
+		PaddingLeft: true,
+	}
+}
+
 func buildLabel(s string) []protocol.InlayHintLabelPart {
 	label := protocol.InlayHintLabelPart{
 		Value: s,

internal/lsp/testdata/inlay_hint/variable_types.go

@@ -0,0 +1,20 @@
+package inlayHint //@inlayHint("package")
+
+func assignTypes() {
+	i, j := 0, len([]string{})-1
+	println(i, j)
+}
+
+func rangeTypes() {
+	for k, v := range []string{} {
+		println(k, v)
+	}
+}
+
+func funcLitType() {
+	myFunc := func(a string) string { return "" }
+}
+
+func compositeLitType() {
+	foo := map[string]interface{}{"": ""}
+}

internal/lsp/testdata/inlay_hint/variable_types.go.golden

@@ -0,0 +1,22 @@
+-- inlayHint --
+package inlayHint //@inlayHint("package")
+
+func assignTypes() {
+	i<int>, j<int> := 0, len([]string{})-1
+	println(i, j)
+}
+
+func rangeTypes() {
+	for k<int>, v<string> := range []string{} {
+		println(k, v)
+	}
+}
+
+func funcLitType() {
+	myFunc<func(a string) string> := func(a string) string { return "" }
+}
+
+func compositeLitType() {
+	foo<map[string]interface{}> := map[string]interface{}{"": ""}
+}
+