gopls/internal/golang: fix folding range for function calls

Make folding ranges for function calls consistent with folding
ranges for function bodies and composite literals:

- keep the closing parenthesis visible
- do not generate a folding range if either the first or last
  argument is on the same line as the opening or closing parenthesis

Also refactor some code, add additional comments and test cases.

Fixes golang/go#70467

gopls/internal/golang/folding_range.go

@@ -11,6 +11,7 @@ import (
 	"sort"
 	"strings"
 
+	"golang.org/x/tools/go/ast/astutil"
 	"golang.org/x/tools/gopls/internal/cache"
 	"golang.org/x/tools/gopls/internal/cache/parsego"
 	"golang.org/x/tools/gopls/internal/file"
@@ -73,6 +74,7 @@ func foldingRangeFunc(pgf *parsego.File, n ast.Node, lineFoldingOnly bool) *Fold
 	// TODO(suzmue): include trailing empty lines before the closing
 	// parenthesis/brace.
 	var kind protocol.FoldingRangeKind
+	// start and end define the range of content to fold away.
 	var start, end token.Pos
 	switch n := n.(type) {
 	case *ast.BlockStmt:
@@ -81,23 +83,41 @@ func foldingRangeFunc(pgf *parsego.File, n ast.Node, lineFoldingOnly bool) *Fold
 		if num := len(n.List); num != 0 {
 			startList, endList = n.List[0].Pos(), n.List[num-1].End()
 		}
-		start, end = validLineFoldingRange(pgf.Tok, n.Lbrace, n.Rbrace, startList, endList, lineFoldingOnly)
+		start, end = getLineFoldingRange(pgf.Tok, n.Lbrace, n.Rbrace, startList, endList, lineFoldingOnly, false)
 	case *ast.CaseClause:
 		// Fold from position of ":" to end.
 		start, end = n.Colon+1, n.End()
 	case *ast.CommClause:
 		// Fold from position of ":" to end.
 		start, end = n.Colon+1, n.End()
 	case *ast.CallExpr:
-		// Fold from position of "(" to position of ")".
-		start, end = n.Lparen+1, n.Rparen
+		// Fold between positions of or lines between "(" and ")".
+		var startArgs, endArgs token.Pos
+		if num := len(n.Args); num != 0 {
+			startArgs, endArgs = n.Args[0].Pos(), n.Args[num-1].End()
+			if n.Ellipsis.IsValid() {
+				endArgs = n.Ellipsis + token.Pos(len("..."))
+			}
+		}
+		start, end = getLineFoldingRange(pgf.Tok, n.Lparen, n.Rparen, startArgs, endArgs, lineFoldingOnly, true)
 	case *ast.FieldList:
 		// Fold between positions of or lines between opening parenthesis/brace and closing parenthesis/brace.
-		var startList, endList token.Pos
+		var (
+			startList, endList   token.Pos
+			canHaveTrailingComma bool
+		)
 		if num := len(n.List); num != 0 {
 			startList, endList = n.List[0].Pos(), n.List[num-1].End()
+
+			path, _ := astutil.PathEnclosingInterval(pgf.File, n.Pos(), n.End())
+			for _, p := range path {
+				if _, ok := p.(*ast.FuncType); ok {
+					canHaveTrailingComma = true
+					break
+				}
+			}
 		}
-		start, end = validLineFoldingRange(pgf.Tok, n.Opening, n.Closing, startList, endList, lineFoldingOnly)
+		start, end = getLineFoldingRange(pgf.Tok, n.Opening, n.Closing, startList, endList, lineFoldingOnly, canHaveTrailingComma)
 	case *ast.GenDecl:
 		// If this is an import declaration, set the kind to be protocol.Imports.
 		if n.Tok == token.IMPORT {
@@ -108,7 +128,7 @@ func foldingRangeFunc(pgf *parsego.File, n ast.Node, lineFoldingOnly bool) *Fold
 		if num := len(n.Specs); num != 0 {
 			startSpecs, endSpecs = n.Specs[0].Pos(), n.Specs[num-1].End()
 		}
-		start, end = validLineFoldingRange(pgf.Tok, n.Lparen, n.Rparen, startSpecs, endSpecs, lineFoldingOnly)
+		start, end = getLineFoldingRange(pgf.Tok, n.Lparen, n.Rparen, startSpecs, endSpecs, lineFoldingOnly, false)
 	case *ast.BasicLit:
 		// Fold raw string literals from position of "`" to position of "`".
 		if n.Kind == token.STRING && len(n.Value) >= 2 && n.Value[0] == '`' && n.Value[len(n.Value)-1] == '`' {
@@ -120,13 +140,17 @@ func foldingRangeFunc(pgf *parsego.File, n ast.Node, lineFoldingOnly bool) *Fold
 		if num := len(n.Elts); num != 0 {
 			startElts, endElts = n.Elts[0].Pos(), n.Elts[num-1].End()
 		}
-		start, end = validLineFoldingRange(pgf.Tok, n.Lbrace, n.Rbrace, startElts, endElts, lineFoldingOnly)
+		start, end = getLineFoldingRange(pgf.Tok, n.Lbrace, n.Rbrace, startElts, endElts, lineFoldingOnly, true)
 	}
 
 	// Check that folding positions are valid.
 	if !start.IsValid() || !end.IsValid() {
 		return nil
 	}
+	if start == end {
+		// Nothing to fold.
+		return nil
+	}
 	// in line folding mode, do not fold if the start and end lines are the same.
 	if lineFoldingOnly && safetoken.Line(pgf.Tok, start) == safetoken.Line(pgf.Tok, end) {
 		return nil
@@ -141,26 +165,49 @@ func foldingRangeFunc(pgf *parsego.File, n ast.Node, lineFoldingOnly bool) *Fold
 	}
 }
 
-// validLineFoldingRange returns start and end token.Pos for folding range if the range is valid.
-// returns token.NoPos otherwise, which fails token.IsValid check
-func validLineFoldingRange(tokFile *token.File, open, close, start, end token.Pos, lineFoldingOnly bool) (token.Pos, token.Pos) {
-	if lineFoldingOnly {
-		if !open.IsValid() || !close.IsValid() {
-			return token.NoPos, token.NoPos
-		}
+// getLineFoldingRange returns the folding range for nodes with parentheses/braces/brackets
+// that potentially can take up multiple lines.
+func getLineFoldingRange(tokFile *token.File, open, close, start, end token.Pos, lineFoldingOnly, canHaveTrailingComma bool) (token.Pos, token.Pos) {
+	if !open.IsValid() || !close.IsValid() {
+		return token.NoPos, token.NoPos
+	}
 
-		// Don't want to fold if the start/end is on the same line as the open/close
-		// as an example, the example below should *not* fold:
-		// var x = [2]string{"d",
-		// "e" }
-		if safetoken.Line(tokFile, open) == safetoken.Line(tokFile, start) ||
-			safetoken.Line(tokFile, close) == safetoken.Line(tokFile, end) {
-			return token.NoPos, token.NoPos
-		}
+	if !lineFoldingOnly {
+		// Can fold between opening and closing parenthesis/brace
+		// even if they are on the same line.
+		return open + 1, close
+	}
+
+	// Clients with "LineFoldingOnly" set to true can fold only full lines.
+	// So, we return a folding range only when the closing parenthesis/brace
+	// and the end of the last argument/statement/element are on different lines.
+	//
+	// We could skip the check for the opening parenthesis/brace and start of
+	// the first argument/statement/element. For example, the following code
+	//
+	//	var x = []string{"a",
+	//	"b",
+	//	"c" }
+	//
+	// can be folded to
+	//
+	//	var x = []string{"a", ...
+	//	"c" }
+	//
+	// However, this might look confusing. So, check the lines of "open" and
+	// "start" positions as well.
+
+	if safetoken.Line(tokFile, open) == safetoken.Line(tokFile, start) ||
+		safetoken.Line(tokFile, close) == safetoken.Line(tokFile, end) {
+		return token.NoPos, token.NoPos
+	}
 
-		return open + 1, end
+	if canHaveTrailingComma && end.IsValid() {
+		// The last argument/statement/element and the closing parenthesis/brace
+		// are on different lines. So, there should be a comma.
+		end++
 	}
-	return open + 1, close
+	return open + 1, end
 }
 
 // commentsFoldingRange returns the folding ranges for all comment blocks in file.

gopls/internal/test/marker/testdata/foldingrange/a.txt

@@ -6,13 +6,17 @@ package folding //@foldingrange(raw)
 import (
 	"fmt"
 	_ "log"
+	"sort"
+	"time"
 )
 
 import _ "os"
 
 // bar is a function.
 // With a multiline doc comment.
-func bar() string {
+func bar() (
+	string,
+) {
 	/* This is a single line comment */
 	switch {
 	case true:
@@ -76,19 +80,74 @@ func bar() string {
 this string
 is not indented`
 }
+
+func _() {
+	slice := []int{1, 2, 3}
+	sort.Slice(slice, func(i, j int) bool {
+		a, b := slice[i], slice[j]
+		return a < b
+	})
+
+	sort.Slice(slice, func(i, j int) bool { return slice[i] < slice[j] })
+
+	sort.Slice(
+		slice,
+		func(i, j int) bool {
+			return slice[i] < slice[j]
+		},
+	)
+
+	fmt.Println(
+		1, 2, 3,
+		4,
+	)
+
+	fmt.Println(1, 2, 3,
+		4, 5, 6,
+		7, 8, 9,
+		10)
+
+	// Call with ellipsis.
+	_ = fmt.Errorf(
+		"test %d %d",
+		[]any{1, 2, 3}...,
+	)
+
+	// Check multiline string.
+	fmt.Println(
+		`multi
+		line
+		string
+		`,
+		1, 2, 3,
+	)
+
+	// Call without arguments.
+	_ = time.Now()
+}
+
+func _(
+	a int, b int,
+	c int,
+) {
+}
 -- @raw --
 package folding //@foldingrange(raw)
 
 import (<0 kind="imports">
 	"fmt"
 	_ "log"
+	"sort"
+	"time"
 </0>)
 
 import _ "os"
 
 // bar is a function.<1 kind="comment">
 // With a multiline doc comment.</1>
-func bar(<2 kind=""></2>) string {<3 kind="">
+func bar() (<2 kind="">
+	string,
+</2>) {<3 kind="">
 	/* This is a single line comment */
 	switch {<4 kind="">
 	case true:<5 kind="">
@@ -152,3 +211,54 @@ func bar(<2 kind=""></2>) string {<3 kind="">
 this string
 is not indented`</34>
 </3>}
+
+func _() {<35 kind="">
+	slice := []int{<36 kind="">1, 2, 3</36>}
+	sort.Slice(<37 kind="">slice, func(<38 kind="">i, j int</38>) bool {<39 kind="">
+		a, b := slice[i], slice[j]
+		return a < b
+	</39>}</37>)
+
+	sort.Slice(<40 kind="">slice, func(<41 kind="">i, j int</41>) bool {<42 kind=""> return slice[i] < slice[j] </42>}</40>)
+
+	sort.Slice(<43 kind="">
+		slice,
+		func(<44 kind="">i, j int</44>) bool {<45 kind="">
+			return slice[i] < slice[j]
+		</45>},
+	</43>)
+
+	fmt.Println(<46 kind="">
+		1, 2, 3,
+		4,
+	</46>)
+
+	fmt.Println(<47 kind="">1, 2, 3,
+		4, 5, 6,
+		7, 8, 9,
+		10</47>)
+
+	// Call with ellipsis.
+	_ = fmt.Errorf(<48 kind="">
+		"test %d %d",
+		[]any{<49 kind="">1, 2, 3</49>}...,
+	</48>)
+
+	// Check multiline string.
+	fmt.Println(<50 kind="">
+		<51 kind="">`multi
+		line
+		string
+		`</51>,
+		1, 2, 3,
+	</50>)
+
+	// Call without arguments.
+	_ = time.Now()
+</35>}
+
+func _(<52 kind="">
+	a int, b int,
+	c int,
+</52>) {<53 kind="">
+</53>}