gopls/completion: prefer rangeable funcs in range statements

Teach gopls that `range` statements can accept iterator funcs per
upcoming Go 1.23 language change (and Go 1.22 "rangefunc" experiment).

I didn't bother disabling this preference for earlier versions of Go
since false positive completions seem unlikely.

For golang/go#66637

Change-Id: Id57687f3fa205fa8e4b4616eebee471e6d11d802
Reviewed-on: https://go-review.googlesource.com/c/tools/+/592915
Reviewed-by: Robert Findley <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Carlos Amedee <[email protected]>
Run-TryBot: Muir Manders <[email protected]>
Auto-Submit: Robert Findley <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>

Author: muirdm

gopls/internal/golang/completion/completion.go

@@ -2134,6 +2134,9 @@ const (
 	kindError
 	kindStringer
 	kindFunc
+	kindRange0Func
+	kindRange1Func
+	kindRange2Func
 )
 
 // penalizedObj represents an object that should be disfavored as a
@@ -2430,8 +2433,12 @@ Nodes:
 		case *ast.RangeStmt:
 			if goplsastutil.NodeContains(node.X, c.pos) {
 				inf.objKind |= kindSlice | kindArray | kindMap | kindString
-				if node.Value == nil {
-					inf.objKind |= kindChan
+				if node.Key == nil && node.Value == nil {
+					inf.objKind |= kindRange0Func | kindRange1Func | kindRange2Func
+				} else if node.Value == nil {
+					inf.objKind |= kindChan | kindRange1Func | kindRange2Func
+				} else {
+					inf.objKind |= kindRange2Func
 				}
 			}
 			return inf
@@ -3233,15 +3240,17 @@ var (
 
 	// "interface { String() string }" (i.e. fmt.Stringer)
 	stringerIntf = types.NewInterfaceType([]*types.Func{
-		types.NewFunc(token.NoPos, nil, "String", types.NewSignature(
-			nil,
-			nil,
+		types.NewFunc(token.NoPos, nil, "String", types.NewSignatureType(
+			nil, nil,
+			nil, nil,
 			types.NewTuple(types.NewParam(token.NoPos, nil, "", types.Typ[types.String])),
 			false,
 		)),
 	}, nil).Complete()
 
 	byteType = types.Universe.Lookup("byte").Type()
+
+	boolType = types.Universe.Lookup("bool").Type()
 )
 
 // candKind returns the objKind of candType, if any.
@@ -3287,7 +3296,16 @@ func candKind(candType types.Type) objKind {
 			kind |= kindBool
 		}
 	case *types.Signature:
-		return kindFunc
+		kind |= kindFunc
+
+		switch rangeFuncParamCount(t) {
+		case 0:
+			kind |= kindRange0Func
+		case 1:
+			kind |= kindRange1Func
+		case 2:
+			kind |= kindRange2Func
+		}
 	}
 
 	if types.Implements(candType, errorIntf) {
@@ -3301,6 +3319,24 @@ func candKind(candType types.Type) objKind {
 	return kind
 }
 
+// If sig looks like a range func, return param count, else return -1.
+func rangeFuncParamCount(sig *types.Signature) int {
+	if sig.Results().Len() != 0 || sig.Params().Len() != 1 {
+		return -1
+	}
+
+	yieldSig, _ := sig.Params().At(0).Type().Underlying().(*types.Signature)
+	if yieldSig == nil {
+		return -1
+	}
+
+	if yieldSig.Results().Len() != 1 || yieldSig.Results().At(0).Type() != boolType {
+		return -1
+	}
+
+	return yieldSig.Params().Len()
+}
+
 // innermostScope returns the innermost scope for c.pos.
 func (c *completer) innermostScope() *types.Scope {
 	for _, s := range c.scopes {

gopls/internal/test/marker/testdata/completion/range_func.txt

@@ -0,0 +1,23 @@
+This test shows we prefer rangeable funcs in range statements.
+
+-- flags --
+-ignore_extra_diags
+
+-- range_func.go --
+package rangefunc
+
+func iterNot(func(int)) {}
+func iter0(func() bool) {}
+func iter1(func(int) bool) {}
+func iter2(func(int, int) bool)
+
+func _() {
+	for range i { //@rankl(" {", "iter0", "iterNot"),rankl(" {", "iter1", "iterNot"),rankl(" {", "iter2", "iterNot")
+	}
+
+	for k := range i { //@rankl(" {", "iter1", "iterNot"),rankl(" {", "iter1", "iter0"),rankl(" {", "iter2", "iter0")
+	}
+
+	for k, v := range i { //@rankl(" {", "iter2", "iterNot"),rankl(" {", "iter2", "iter0"),rankl(" {", "iter2", "iter1")
+	}
+}