internal/lsp/source: improve completion support for args to builtins

We now understand what "kind" of type is expected when using various
builtins. For example, when completing "close(<>)" we prefer channels,
and when completing "delete(<>)" we prefer maps.

I also added some code to infer the expected type for the second
argument to "delete()" and for the args to "copy()":

delete(map[someType]int{}, <>) // expect "someType"

copy([]int{}, <>) // expect "[]int"
copy(<>, []int{}) // expect "[]int"

And I marked "new()" as expected a type name, and it infers the type
name properly:

var _ *int = new(<>) // expected type at "<>" is "int"

Fixes golang/go#36326.

Change-Id: I4295c8753f8341d47010a0553fd2d0c2586f2efa
Reviewed-on: https://go-review.googlesource.com/c/tools/+/212957
Run-TryBot: Muir Manders <muir@mnd.rs>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

internal/lsp/source/completion.go

@@ -1202,12 +1202,26 @@ const (
 	array                   // make an array type ("[2]" in "[2]int")
 )
 
+type objKind int
+
+const (
+	kindArray objKind = 1 << iota
+	kindSlice
+	kindChan
+	kindMap
+	kindStruct
+	kindString
+)
+
 // typeInference holds information we have inferred about a type that can be
 // used at the current position.
 type typeInference struct {
 	// objType is the desired type of an object used at the query position.
 	objType types.Type
 
+	// objKind is a mask of expected kinds of types such as "map", "slice", etc.
+	objKind objKind
+
 	// variadic is true if objType is a slice type from an initial
 	// variadic param.
 	variadic bool
@@ -1334,30 +1348,26 @@ Nodes:
 				}
 
 				if funIdent, ok := node.Fun.(*ast.Ident); ok {
-					switch c.pkg.GetTypesInfo().ObjectOf(funIdent) {
-					case types.Universe.Lookup("append"):
+					obj := c.pkg.GetTypesInfo().ObjectOf(funIdent)
+
+					if obj != nil && obj.Parent() == types.Universe {
+						inf.objKind |= c.builtinArgKind(obj, node)
+
+						if obj.Name() == "new" {
+							inf.typeName.wantTypeName = true
+						}
+
+						// Defer call to builtinArgType so we can provide it the
+						// inferred type from its parent node.
 						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()
-								}
-							}
+							inf.objType, inf.variadic = c.builtinArgType(obj, node, inf.objType)
 						}()
 
-						// The expected type of append() arguments is the expected
-						// type of the append() call itself. For example:
+						// The expected type of builtin arguments like append() is
+						// the expected type of the builtin call itself. For
+						// example:
 						//
-						// var foo []int
-						// foo = append(<>)
+						// var foo []int = append(<>)
 						//
 						// To find the expected type at <> we "skip" the append()
 						// node and get the expected type one level up, which is
@@ -1663,6 +1673,12 @@ func (c *completer) matchingCandidate(cand *candidate) bool {
 
 	typeMatches := func(expType, candType types.Type) bool {
 		if expType == nil {
+			// If we don't expect a specific type, check if we expect a particular
+			// kind of object (map, slice, etc).
+			if c.expectedType.objKind > 0 {
+				return c.expectedType.objKind&candKind(candType) > 0
+			}
+
 			return false
 		}
 
@@ -1776,3 +1792,29 @@ func (c *completer) matchingTypeName(cand *candidate) bool {
 	// Default to saying any type name is a match.
 	return true
 }
+
+// candKind returns the objKind of candType, if any.
+func candKind(candType types.Type) objKind {
+	switch t := candType.Underlying().(type) {
+	case *types.Array:
+		return kindArray
+	case *types.Slice:
+		return kindSlice
+	case *types.Chan:
+		return kindChan
+	case *types.Map:
+		return kindMap
+	case *types.Pointer:
+		// Some builtins handle array pointers as arrays, so just report a pointer
+		// to an array as an array.
+		if _, isArray := t.Elem().Underlying().(*types.Array); isArray {
+			return kindArray
+		}
+	case *types.Basic:
+		if t.Info()&types.IsString > 0 {
+			return kindString
+		}
+	}
+
+	return 0
+}

internal/lsp/source/completion_builtin.go

@@ -0,0 +1,110 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package source
+
+import (
+	"go/ast"
+	"go/types"
+)
+
+// builtinArgKind determines the expected object kind for a builtin
+// argument. It attempts to use the AST hints from builtin.go where
+// possible.
+func (c *completer) builtinArgKind(obj types.Object, call *ast.CallExpr) objKind {
+	astObj := c.snapshot.View().BuiltinPackage().Lookup(obj.Name())
+	if astObj == nil {
+		return 0
+	}
+
+	exprIdx := indexExprAtPos(c.pos, call.Args)
+
+	decl, ok := astObj.Decl.(*ast.FuncDecl)
+	if !ok || exprIdx >= len(decl.Type.Params.List) {
+		return 0
+	}
+
+	switch ptyp := decl.Type.Params.List[exprIdx].Type.(type) {
+	case *ast.ChanType:
+		return kindChan
+	case *ast.ArrayType:
+		return kindSlice
+	case *ast.MapType:
+		return kindMap
+	case *ast.Ident:
+		switch ptyp.Name {
+		case "Type":
+			switch obj.Name() {
+			case "make":
+				return kindChan | kindSlice | kindMap
+			case "len":
+				return kindSlice | kindMap | kindArray | kindString | kindChan
+			case "cap":
+				return kindSlice | kindArray | kindChan
+			}
+		}
+	}
+
+	return 0
+}
+
+// builtinArgType infers the type of an argument to a builtin
+// function. "parentType" is the inferred type for the builtin call's
+// parent node.
+func (c *completer) builtinArgType(obj types.Object, call *ast.CallExpr, parentType types.Type) (infType types.Type, variadic bool) {
+	exprIdx := indexExprAtPos(c.pos, call.Args)
+
+	switch obj.Name() {
+	case "append":
+		// Check if we are completing the variadic append() param.
+		variadic = exprIdx == 1 && len(call.Args) <= 2
+		infType = parentType
+
+		// If we are completing an individual element of the variadic
+		// param, "deslice" the expected type.
+		if !variadic && exprIdx > 0 {
+			if slice, ok := parentType.(*types.Slice); ok {
+				infType = slice.Elem()
+			}
+		}
+	case "delete":
+		if exprIdx > 0 && len(call.Args) > 0 {
+			// Try to fill in expected type of map key.
+			firstArgType := c.pkg.GetTypesInfo().TypeOf(call.Args[0])
+			if firstArgType != nil {
+				if mt, ok := firstArgType.Underlying().(*types.Map); ok {
+					infType = mt.Key()
+				}
+			}
+		}
+	case "copy":
+		var t1, t2 types.Type
+		if len(call.Args) > 0 {
+			t1 = c.pkg.GetTypesInfo().TypeOf(call.Args[0])
+			if len(call.Args) > 1 {
+				t2 = c.pkg.GetTypesInfo().TypeOf(call.Args[1])
+			}
+		}
+
+		// Fill in expected type of either arg if the other is already present.
+		if exprIdx == 1 && t1 != nil {
+			infType = t1
+		} else if exprIdx == 0 && t2 != nil {
+			infType = t2
+		}
+	case "new":
+		if parentType != nil {
+			// Expected type for "new" is the de-pointered parent type.
+			if ptr, ok := parentType.Underlying().(*types.Pointer); ok {
+				infType = ptr.Elem()
+			}
+		}
+	case "make":
+		if exprIdx == 0 {
+			infType = parentType
+		}
+	}
+
+	return infType, variadic
+}

internal/lsp/testdata/builtins/builtin_args.go

@@ -0,0 +1,37 @@
+package builtins
+
+func _() {
+	var (
+		slice_    []int          //@item(builtinSlice, "slice_", "[]int", "var")
+		map_      map[string]int //@item(builtinMap, "map_", "map[string]int", "var")
+		string_   string         //@item(builtinString, "string_", "string", "var")
+		array_    [0]int         //@item(builtinArray, "array_", "[0]int", "var")
+		arrayPtr_ *[0]int        //@item(builtinArrayPtr, "arrayPtr_", "*[0]int", "var")
+		chan_     chan int       //@item(builtinChan, "chan_", "chan int", "var")
+		ptr_      *int           //@item(builtinPtr, "ptr_", "*int", "var")
+		int_      int            //@item(builtinInt, "int_", "int", "var")
+	)
+
+	close() //@rank(")", builtinChan, builtinSlice)
+
+	append() //@rank(")", builtinSlice, builtinChan)
+
+	copy()          //@rank(")", builtinSlice, builtinChan)
+	copy(slice_, s) //@rank(")", builtinSlice, builtinString)
+	copy(s, slice_) //@rank(",", builtinSlice, builtinString)
+
+	delete()        //@rank(")", builtinMap, builtinChan)
+	delete(map_, s) //@rank(")", builtinString, builtinSlice)
+
+	len() //@rank(")", builtinSlice, builtinInt),rank(")", builtinMap, builtinInt),rank(")", builtinString, builtinInt),rank(")", builtinArray, builtinInt),rank(")", builtinArrayPtr, builtinPtr),rank(")", builtinChan, builtinInt)
+
+	cap() //@rank(")", builtinSlice, builtinMap),rank(")", builtinArray, builtinString),rank(")", builtinArrayPtr, builtinPtr),rank(")", builtinChan, builtinInt)
+
+	make() //@rank(")", builtinMap, builtinInt),rank(")", builtinChan, builtinInt),rank(")", builtinSlice, builtinInt)
+
+	var _ []int = make() //@rank(")", builtinSlice, builtinMap)
+
+	type myStruct struct{}  //@item(builtinStructType, "myStruct", "struct{...}", "struct")
+	new()                   //@rank(")", builtinStructType, builtinInt)
+	var _ *myStruct = new() //@rank(")", builtinStructType, int)
+}

internal/lsp/testdata/summary.txt.golden

@@ -4,7 +4,7 @@ CompletionSnippetCount = 62
 UnimportedCompletionsCount = 4
 DeepCompletionsCount = 5
 FuzzyCompletionsCount = 8
-RankedCompletionsCount = 33
+RankedCompletionsCount = 56
 CaseSensitiveCompletionsCount = 4
 DiagnosticsCount = 35
 FoldingRangesCount = 2