From e0b9cf74f60215fbb96eab9fc3c7e2ee2de2df71 Mon Sep 17 00:00:00 2001
From: Muir Manders <muir@mnd.rs>
Date: Fri, 28 May 2021 14:59:19 -0700
Subject: [PATCH] lsp/completion: support completing to Elem() types

For array, slice, maps and chan candidates, we now support
transforming them to their element type in completions. For example:

    var m map[string]int
    var _ int = m<>

At <> we complete to "m[]" because we see that the map value type
matches our expected type.

Fixes golang/go#46045.

Change-Id: Ibee088550193a53744f93217cc365f67f301ae90
Reviewed-on: https://go-review.googlesource.com/c/tools/+/323451
Run-TryBot: Muir Manders <muir@mnd.rs>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Trust: Cherry Mui <cherryyz@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>
---
 internal/lsp/source/completion/builtin.go        |  3 +++
 internal/lsp/source/completion/completion.go     | 16 ++++++++++++++++
 .../lsp/source/completion/deep_completion.go     |  6 ++++++
 internal/lsp/source/completion/format.go         |  4 ++++
 internal/lsp/testdata/summary.txt.golden         |  2 +-
 internal/lsp/testdata/typemods/type_mods.go      |  6 ++++++
 6 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/internal/lsp/source/completion/builtin.go b/internal/lsp/source/completion/builtin.go
index 086e3f788d..39732d8643 100644
--- a/internal/lsp/source/completion/builtin.go
+++ b/internal/lsp/source/completion/builtin.go
@@ -73,6 +73,9 @@ func (c *completer) builtinArgType(obj types.Object, call *ast.CallExpr, parentI
 			// Infer first append() arg type as apparent return type of
 			// append().
 			inf.objType = parentInf.objType
+			if parentInf.variadic {
+				inf.objType = types.NewSlice(inf.objType)
+			}
 			break
 		}
 
diff --git a/internal/lsp/source/completion/completion.go b/internal/lsp/source/completion/completion.go
index 14a2052953..741e6b35b6 100644
--- a/internal/lsp/source/completion/completion.go
+++ b/internal/lsp/source/completion/completion.go
@@ -1777,6 +1777,7 @@ const (
 	invoke                           // make a function call: "()" in "foo()"
 	takeSlice                        // take slice of array: "[:]" in "foo[:]"
 	takeDotDotDot                    // turn slice into variadic args: "..." in "foo..."
+	index                            // index into slice/array: "[0]" in "foo[0]"
 )
 
 type objKind int
@@ -2404,6 +2405,9 @@ func derivableTypes(t types.Type, addressable bool, f func(t types.Type, address
 			return true
 		}
 	case *types.Array:
+		if f(t.Elem(), true, index) {
+			return true
+		}
 		// Try converting array to slice.
 		if f(types.NewSlice(t.Elem()), false, takeSlice) {
 			return true
@@ -2412,6 +2416,18 @@ func derivableTypes(t types.Type, addressable bool, f func(t types.Type, address
 		if f(t.Elem(), false, dereference) {
 			return true
 		}
+	case *types.Slice:
+		if f(t.Elem(), true, index) {
+			return true
+		}
+	case *types.Map:
+		if f(t.Elem(), false, index) {
+			return true
+		}
+	case *types.Chan:
+		if f(t.Elem(), false, chanRead) {
+			return true
+		}
 	}
 
 	// Check if c is addressable and a pointer to c matches our type inference.
diff --git a/internal/lsp/source/completion/deep_completion.go b/internal/lsp/source/completion/deep_completion.go
index e7c4854e23..45a02ff071 100644
--- a/internal/lsp/source/completion/deep_completion.go
+++ b/internal/lsp/source/completion/deep_completion.go
@@ -260,6 +260,12 @@ func (c *completer) addCandidate(ctx context.Context, cand *candidate) {
 		cand.score *= 1.1
 	}
 
+	// Slight penalty for index modifier (e.g. changing "foo" to
+	// "foo[]") to curb false positives.
+	if cand.hasMod(index) {
+		cand.score *= 0.9
+	}
+
 	// Favor shallow matches by lowering score according to depth.
 	cand.score -= cand.score * c.deepState.scorePenalty(cand)
 
diff --git a/internal/lsp/source/completion/format.go b/internal/lsp/source/completion/format.go
index 08b933f444..5a20633f7c 100644
--- a/internal/lsp/source/completion/format.go
+++ b/internal/lsp/source/completion/format.go
@@ -132,6 +132,10 @@ Suffixes:
 			suffix += "[:]"
 		case takeDotDotDot:
 			suffix += "..."
+		case index:
+			snip.WriteText("[")
+			snip.WritePlaceholder(nil)
+			snip.WriteText("]")
 		}
 	}
 
diff --git a/internal/lsp/testdata/summary.txt.golden b/internal/lsp/testdata/summary.txt.golden
index 1c1b780157..ce4d7e4acf 100644
--- a/internal/lsp/testdata/summary.txt.golden
+++ b/internal/lsp/testdata/summary.txt.golden
@@ -2,7 +2,7 @@
 CallHierarchyCount = 2
 CodeLensCount = 5
 CompletionsCount = 265
-CompletionSnippetCount = 102
+CompletionSnippetCount = 103
 UnimportedCompletionsCount = 5
 DeepCompletionsCount = 5
 FuzzyCompletionsCount = 8
diff --git a/internal/lsp/testdata/typemods/type_mods.go b/internal/lsp/testdata/typemods/type_mods.go
index 479b481cbe..f5f0f80767 100644
--- a/internal/lsp/testdata/typemods/type_mods.go
+++ b/internal/lsp/testdata/typemods/type_mods.go
@@ -13,3 +13,9 @@ func fooPtr() *int { //@item(modFooPtr, "fooPtr", "func() *int", "func")
 func _() {
 	var _ int = foo //@snippet(" //", modFooFunc, "fooFunc()()", "fooFunc()()"),snippet(" //", modFooPtr, "*fooPtr()", "*fooPtr()")
 }
+
+func _() {
+	var m map[int][]chan int //@item(modMapChanPtr, "m", "map[int]chan *int", "var")
+
+	var _ int = m //@snippet(" //", modMapChanPtr, "<-m[${1:}][${2:}]", "<-m[${1:}][${2:}]")
+}