From 97ad0ed33101b39103eaf8aeab62b59dd1cfb89e Mon Sep 17 00:00:00 2001
From: Muir Manders <muir@mnd.rs>
Date: Wed, 6 Nov 2019 11:39:30 -0800
Subject: [PATCH] internal/lsp: improve completion in append() calls

Add a special case for append() arguments so we infer the expected
type from the append() context. For example:

var foo []int
foo = append(<>)

We now infer the expected type at <> to be []int. We also support the
variadicity of append().

Change-Id: Ie0ef0007907fcb7992f9697cb90970ce4d9a66b8
Reviewed-on: https://go-review.googlesource.com/c/tools/+/205606
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
---
 internal/lsp/source/completion.go        | 39 +++++++++++++++++++++---
 internal/lsp/testdata/append/append.go   | 19 ++++++++++++
 internal/lsp/testdata/summary.txt.golden |  4 +--
 3 files changed, 56 insertions(+), 6 deletions(-)
 create mode 100644 internal/lsp/testdata/append/append.go

diff --git a/internal/lsp/source/completion.go b/internal/lsp/source/completion.go
index edcffa56ea..277d2cee89 100644
--- a/internal/lsp/source/completion.go
+++ b/internal/lsp/source/completion.go
@@ -1053,10 +1053,8 @@ type typeNameInference struct {
 
 // expectedType returns information about the expected type for an expression at
 // the query position.
-func expectedType(c *completer) typeInference {
-	inf := typeInference{
-		typeName: expectTypeName(c),
-	}
+func expectedType(c *completer) (inf typeInference) {
+	inf.typeName = expectTypeName(c)
 
 	if c.enclosingCompositeLiteral != nil {
 		inf.objType = c.expectedCompositeLiteralType()
@@ -1139,6 +1137,39 @@ Nodes:
 						break Nodes
 					}
 				}
+
+				if funIdent, ok := node.Fun.(*ast.Ident); ok {
+					switch c.pkg.GetTypesInfo().ObjectOf(funIdent) {
+					case types.Universe.Lookup("append"):
+						defer func() {
+							exprIdx := indexExprAtPos(c.pos, node.Args)
+
+							// Check if we are completing the variadic append()
+							// param. We defer this since we don't want to inherit
+							// variadicity from the next node.
+							inf.variadic = exprIdx == 1 && len(node.Args) <= 2
+
+							// If we are completing an individual element of the
+							// variadic param, "deslice" the expected type.
+							if !inf.variadic && exprIdx > 0 {
+								if slice, ok := inf.objType.(*types.Slice); ok {
+									inf.objType = slice.Elem()
+								}
+							}
+						}()
+
+						// The expected type of append() arguments is the expected
+						// type of the append() call itself. For example:
+						//
+						// var foo []int
+						// foo = append(<>)
+						//
+						// To find the expected type at <> we "skip" the append()
+						// node and get the expected type one level up, which is
+						// []int.
+						continue Nodes
+					}
+				}
 			}
 			return inf
 		case *ast.ReturnStmt:
diff --git a/internal/lsp/testdata/append/append.go b/internal/lsp/testdata/append/append.go
new file mode 100644
index 0000000000..6c6bfb797e
--- /dev/null
+++ b/internal/lsp/testdata/append/append.go
@@ -0,0 +1,19 @@
+package append
+
+func foo([]string)  {}
+func bar(...string) {}
+
+func _() {
+	var (
+		aInt     []int    //@item(appendInt, "aInt", "[]int", "var")
+		aStrings []string //@item(appendStrings, "aStrings", "[]string", "var")
+		aString  string   //@item(appendString, "aString", "string", "var")
+	)
+
+	foo(append())           //@rank("))", appendStrings, appendInt),rank("))", appendStrings, appendString)
+	foo(append(nil, a))     //@rank("))", appendStrings, appendInt),rank("))", appendString, appendInt),snippet("))", appendStrings, "aStrings...", "aStrings...")
+	foo(append(nil, "", a)) //@rank("))", appendString, appendInt),rank("))", appendString, appendStrings)
+
+	// Don't add "..." to append() argument.
+	bar(append()) //@snippet("))", appendStrings, "aStrings", "aStrings")
+}
diff --git a/internal/lsp/testdata/summary.txt.golden b/internal/lsp/testdata/summary.txt.golden
index a875e5a8ed..aacef0be0b 100644
--- a/internal/lsp/testdata/summary.txt.golden
+++ b/internal/lsp/testdata/summary.txt.golden
@@ -1,10 +1,10 @@
 -- summary --
 CompletionsCount = 215
-CompletionSnippetCount = 45
+CompletionSnippetCount = 47
 UnimportedCompletionsCount = 3
 DeepCompletionsCount = 5
 FuzzyCompletionsCount = 7
-RankedCompletionsCount = 16
+RankedCompletionsCount = 22
 CaseSensitiveCompletionsCount = 4
 DiagnosticsCount = 22
 FoldingRangesCount = 2