internal/lsp/source: improve completion involving multiple return values

For example:

// Prefer functions that return one or two values. Previously
// we had no preference.
foo, bar := <>

// Prefer functions that return "(int)" or "(int, ??)". Previously we
// only preferred the former.
var foo int
foo, bar := <>

// Prefer functions that return "(int)" or "(int, int)". Previously we
// only preferred the former.
var foo func(int, int)
foo(<>)

In the above example, we don't handle "foo" being variadic yet.

I also took the liberty to break up matchingCandidate() into separate
functions since it was getting rather long.

Updates golang/go#36540.

Change-Id: I9140dd989dfde1ddcfcd9d2a14198045c02587f2
Reviewed-on: https://go-review.googlesource.com/c/tools/+/215537
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

@@ -1270,6 +1270,15 @@ type candidateInference struct {
 	// typeName holds information about the expected type name at
 	// position, if any.
 	typeName typeNameInference
+
+	// assignees are the types that would receive a function call's
+	// results at the position. For example:
+	//
+	// foo := 123
+	// foo, bar := <>
+	//
+	// at "<>", the assignees are [int, <invalid>].
+	assignees []types.Type
 }
 
 // typeNameInference holds information about the expected type name at
@@ -1322,6 +1331,20 @@ Nodes:
 				if tv, ok := c.pkg.GetTypesInfo().Types[node.Lhs[i]]; ok {
 					inf.objType = tv.Type
 				}
+
+				// If we have a single expression on the RHS, record the LHS
+				// assignees so we can favor multi-return function calls with
+				// matching result values.
+				if len(node.Rhs) <= 1 {
+					for _, lhs := range node.Lhs {
+						inf.assignees = append(inf.assignees, c.pkg.GetTypesInfo().TypeOf(lhs))
+					}
+				} else {
+					// Otherwse, record our single assignee, even if its type is
+					// not available. We use this info to downrank functions
+					// with the wrong number of result values.
+					inf.assignees = append(inf.assignees, c.pkg.GetTypesInfo().TypeOf(node.Lhs[i]))
+				}
 			}
 			return inf
 		case *ast.ValueSpec:
@@ -1352,6 +1375,17 @@ Nodes:
 							beyondLastParam = exprIdx >= numParams
 						)
 
+						// If we have one or zero arg expressions, we may be
+						// completing to a function call that returns multiple
+						// values, in turn getting passed in to the surrounding
+						// call. Record the assignees so we can favor function
+						// calls that return matching values.
+						if len(node.Args) <= 1 {
+							for i := 0; i < sig.Params().Len(); i++ {
+								inf.assignees = append(inf.assignees, sig.Params().At(i).Type())
+							}
+						}
+
 						if sig.Variadic() {
 							// If we are beyond the last param or we are the last
 							// param w/ further expressions, we expect a single
@@ -1701,61 +1735,17 @@ func (c *completer) matchingCandidate(cand *candidate) bool {
 	// are invoked by default.
 	cand.expandFuncCall = isFunc(cand.obj)
 
-	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.inference.objKind > 0 {
-				return c.inference.objKind&candKind(candType) > 0
-			}
-
-			return false
-		}
-
-		// Take into account any type modifiers on the expected type.
-		candType = c.inference.applyTypeModifiers(candType, cand.addressable)
-		if candType == nil {
-			return false
-		}
-
-		// Handle untyped values specially since AssignableTo gives false negatives
-		// for them (see https://golang.org/issue/32146).
-		if candBasic, ok := candType.Underlying().(*types.Basic); ok {
-			if wantBasic, ok := expType.Underlying().(*types.Basic); ok {
-				// Make sure at least one of them is untyped.
-				if isUntyped(candType) || isUntyped(expType) {
-					// Check that their constant kind (bool|int|float|complex|string) matches.
-					// This doesn't take into account the constant value, so there will be some
-					// false positives due to integer sign and overflow.
-					if candBasic.Info()&types.IsConstType == wantBasic.Info()&types.IsConstType {
-						// Lower candidate score if the types are not identical. This avoids
-						// ranking untyped constants above candidates with an exact type
-						// match. Don't lower score of builtin constants (e.g. "true").
-						if !types.Identical(candType, expType) && cand.obj.Parent() != types.Universe {
-							cand.score /= 2
-						}
-						return true
-					}
-				}
-			}
-		}
-
-		// AssignableTo covers the case where the types are equal, but also handles
-		// cases like assigning a concrete type to an interface type.
-		return types.AssignableTo(candType, expType)
-	}
-
-	if typeMatches(c.inference.objType, candType) {
+	if c.inference.typeMatches(cand, c.inference.objType, candType) {
 		// If obj's type matches, we don't want to expand to an invocation of obj.
 		cand.expandFuncCall = false
 		return true
 	}
 
 	// Try using a function's return type as its type.
-	if sig, ok := candType.Underlying().(*types.Signature); ok && sig.Results().Len() == 1 {
-		if typeMatches(c.inference.objType, sig.Results().At(0).Type()) {
-			// If obj's return value matches the expected type, we need to invoke obj
-			// in the completion.
+	if sig, ok := candType.Underlying().(*types.Signature); ok {
+		if c.inference.signatureMatches(cand, sig) {
+			// If obj's signature's return value matches the expected type,
+			// we need to invoke obj in the completion.
 			cand.expandFuncCall = true
 			return true
 		}
@@ -1764,8 +1754,10 @@ func (c *completer) matchingCandidate(cand *candidate) bool {
 	// When completing the variadic parameter, if the expected type is
 	// []T then check candType against T.
 	if c.inference.variadic {
-		if slice, ok := c.inference.objType.(*types.Slice); ok && typeMatches(slice.Elem(), candType) {
-			return true
+		if slice, ok := c.inference.objType.(*types.Slice); ok {
+			if c.inference.typeMatches(cand, slice.Elem(), candType) {
+				return true
+			}
 		}
 	}
 
@@ -1792,6 +1784,97 @@ func (c *completer) matchingCandidate(cand *candidate) bool {
 	return false
 }
 
+// typeMatches reports whether an object of candType makes a good
+// completion candidate given the expected type expType. The
+// candidate's score may be mutated to downrank the candidate in
+// certain situations.
+func (ci *candidateInference) typeMatches(cand *candidate, 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 ci.objKind > 0 {
+			return ci.objKind&candKind(candType) > 0
+		}
+
+		return false
+	}
+
+	// Take into account any type modifiers on the expected type.
+	candType = ci.applyTypeModifiers(candType, cand.addressable)
+	if candType == nil {
+		return false
+	}
+
+	// Handle untyped values specially since AssignableTo gives false negatives
+	// for them (see https://golang.org/issue/32146).
+	if candBasic, ok := candType.Underlying().(*types.Basic); ok {
+		if wantBasic, ok := expType.Underlying().(*types.Basic); ok {
+			// Make sure at least one of them is untyped.
+			if isUntyped(candType) || isUntyped(expType) {
+				// Check that their constant kind (bool|int|float|complex|string) matches.
+				// This doesn't take into account the constant value, so there will be some
+				// false positives due to integer sign and overflow.
+				if candBasic.Info()&types.IsConstType == wantBasic.Info()&types.IsConstType {
+					// Lower candidate score if the types are not identical. This avoids
+					// ranking untyped constants above candidates with an exact type
+					// match. Don't lower score of builtin constants (e.g. "true").
+					if !types.Identical(candType, expType) && cand.obj.Parent() != types.Universe {
+						cand.score /= 2
+					}
+					return true
+				}
+			}
+		}
+	}
+
+	// AssignableTo covers the case where the types are equal, but also handles
+	// cases like assigning a concrete type to an interface type.
+	return types.AssignableTo(candType, expType)
+}
+
+// signatureMatches reports whether an invocation of sig makes a good
+// completion candidate. The candidate's score may be mutated to
+// downrank the candidate in certain situations.
+func (ci *candidateInference) signatureMatches(cand *candidate, sig *types.Signature) bool {
+	// If sig returns a single value and it matches our expected type,
+	// invocation of sig is a good candidate.
+	if sig.Results().Len() == 1 {
+		return ci.typeMatches(cand, ci.objType, sig.Results().At(0).Type())
+	}
+
+	if len(ci.assignees) == 0 {
+		return false
+	}
+
+	// If our signature doesn't return the right number of values, it's
+	// not a match, so downrank it. For example:
+	//
+	//  var foo func() (int, int)
+	//  a, b, c := <> // downrank "foo()" since it only returns two values
+	//
+	// TODO: handle the case when we are completing the parameters to a
+	//       variadic function call.
+	if sig.Results().Len() != len(ci.assignees) {
+		cand.score /= 2
+		return false
+	}
+
+	// If at least one assignee has a valid type, and all valid
+	// assignees match the corresponding sig result value, the signature
+	// is a match.
+	allMatch := false
+	for i, a := range ci.assignees {
+		if a == nil || a.Underlying() == types.Typ[types.Invalid] {
+			continue
+		}
+		allMatch = ci.typeMatches(cand, a, sig.Results().At(i).Type())
+		if !allMatch {
+			break
+		}
+	}
+	return allMatch
+}
+
 func (c *completer) matchingTypeName(cand *candidate) bool {
 	if !c.wantTypeName() {
 		return false

internal/lsp/source/deep_completion.go

@@ -138,7 +138,7 @@ func (c *completer) shouldPrune() bool {
 	return false
 }
 
-// deepSearch searches through obj's subordinate objects for more
+// deepSearch searches through cand's subordinate objects for more
 // completion items.
 func (c *completer) deepSearch(cand candidate) {
 	if c.deepState.maxDepth == 0 {

internal/lsp/testdata/multireturn/multi_return.go

@@ -0,0 +1,33 @@
+package multireturn
+
+func f0() {} //@item(multiF0, "f0", "func()", "func")
+
+func f1(int) int { return 0 } //@item(multiF1, "f1", "func(int) int", "func")
+
+func f2(int, int) (int, int) { return 0, 0 } //@item(multiF2, "f2", "func(int, int) (int, int)", "func")
+
+func f2Str(string, string) (string, string) { return "", "" } //@item(multiF2Str, "f2Str", "func(string, string) (string, string)", "func")
+
+func f3(int, int, int) (int, int, int) { return 0, 0, 0 } //@item(multiF3, "f3", "func(int, int, int) (int, int, int)", "func")
+
+func _() {
+	_ := f //@rank(" //", multiF1, multiF2)
+
+	_, _ := f //@rank(" //", multiF2, multiF0),rank(" //", multiF1, multiF0)
+
+	_, _ := _, f //@rank(" //", multiF1, multiF2),rank(" //", multiF1, multiF0)
+
+	_, _ := f, abc //@rank(", abc", multiF1, multiF2)
+
+	f1()     //@rank(")", multiF1, multiF0)
+	f1(f)    //@rank(")", multiF1, multiF2)
+	f2(f)    //@rank(")", multiF2, multiF3),rank(")", multiF1, multiF3)
+	f2(1, f) //@rank(")", multiF1, multiF2),rank(")", multiF1, multiF0)
+	f2Str()  //@rank(")", multiF2Str, multiF2)
+
+	var i int
+	i, _ := f //@rank(" //", multiF2, multiF2Str)
+
+	var s string
+	_, s := f //@rank(" //", multiF2Str, multiF2)
+}

internal/lsp/testdata/summary.txt.golden

@@ -4,7 +4,7 @@ CompletionSnippetCount = 66
 UnimportedCompletionsCount = 9
 DeepCompletionsCount = 5
 FuzzyCompletionsCount = 8
-RankedCompletionsCount = 68
+RankedCompletionsCount = 83
 CaseSensitiveCompletionsCount = 4
 DiagnosticsCount = 38
 FoldingRangesCount = 2