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,7 +83,7 @@ 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()
@@ -89,15 +91,33 @@ func foldingRangeFunc(pgf *parsego.File, n ast.Node, lineFoldingOnly bool) *Fold
 		// 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
-		}
-
-		// 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
-		}
-
-		return open + 1, end
+// 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
 	}
-	return open + 1, close
+
+	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
+	}
+
+	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, 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>}

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

@@ -15,6 +15,8 @@ package folding //@foldingrange(raw)
 import (
 	"fmt"
 	_ "log"
+	"sort"
+	"time"
 )
 
 import _ "os"
@@ -85,12 +87,65 @@ 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"</0>
+	_ "log"
+	"sort"
+	"time"</0>
 )
 
 import _ "os"
@@ -122,7 +177,7 @@ func bar() string {<2 kind="">
 	_ = []int{<11 kind="">
 		1,
 		2,
-		3</11>,
+		3,</11>
 	}
 	_ = [2]string{"d",
 		"e",
@@ -130,7 +185,7 @@ func bar() string {<2 kind="">
 	_ = map[string]int{<12 kind="">
 		"a": 1,
 		"b": 2,
-		"c": 3</12>,
+		"c": 3,</12>
 	}
 	type T struct {<13 kind="">
 		f string
@@ -140,7 +195,7 @@ func bar() string {<2 kind="">
 	_ = T{<14 kind="">
 		f: "j",
 		g: 4,
-		h: "i"</14>,
+		h: "i",</14>
 	}
 	x, y := make(chan bool), make(chan bool)
 	select {<15 kind="">
@@ -161,3 +216,54 @@ func bar() string {<2 kind="">
 this string
 is not indented`</2></22>
 }
+
+func _() {<23 kind="">
+	slice := []int{1, 2, 3}
+	sort.Slice(slice, func(i, j int) bool {<24 kind="">
+		a, b := slice[i], slice[j]
+		return a < b</24>
+	})
+
+	sort.Slice(slice, func(i, j int) bool { return slice[i] < slice[j] })
+
+	sort.Slice(<25 kind="">
+		slice,
+		func(i, j int) bool {<26 kind="">
+			return slice[i] < slice[j]</26>
+		},</25>
+	)
+
+	fmt.Println(<27 kind="">
+		1, 2, 3,
+		4,</27>
+	)
+
+	fmt.Println(1, 2, 3,
+		4, 5, 6,
+		7, 8, 9,
+		10)
+
+	// Call with ellipsis.
+	_ = fmt.Errorf(<28 kind="">
+		"test %d %d",
+		[]any{1, 2, 3}...,</28>
+	)
+
+	// Check multiline string.
+	fmt.Println(<29 kind="">
+		<30 kind="">`multi
+		line
+		string
+		`</30>,
+		1, 2, 3,</29>
+	)
+
+	// Call without arguments.
+	_ = time.Now()</23>
+}
+
+func _(<31 kind="">
+	a int, b int,
+	c int,</31>
+) {
+}

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

@@ -31,11 +31,11 @@ import (<1 kind="imports">
 	_ "os" </1>)
 	
 // badBar is a function.
-func badBar(<2 kind=""></2>) string {<3 kind=""> x := true
-	if x {<4 kind=""> 
+func badBar() string {<2 kind=""> x := true
+	if x {<3 kind=""> 
 		// This is the only foldable thing in this file when lineFoldingOnly
-		fmt.Println(<5 kind="">"true"</5>)
-	</4>} else {<6 kind="">
-		fmt.Println(<7 kind="">"false"</7>) </6>}
+		fmt.Println(<4 kind="">"true"</4>)
+	</3>} else {<5 kind="">
+		fmt.Println(<6 kind="">"false"</6>) </5>}
 	return ""
-</3>}
+</2>}

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

@@ -0,0 +1,26 @@
+This test verifies that textDocument/foldingRange does not panic
+and produces no folding ranges if a file contains errors.
+
+-- flags --
+-ignore_extra_diags
+
+-- a.go --
+package folding //@foldingrange(raw)
+
+// No comma.
+func _(
+	a string
+) {}
+
+// Extra brace.
+func _() {}}
+-- @raw --
+package folding //@foldingrange(raw)
+
+// No comma.
+func _(
+	a string
+) {}
+
+// Extra brace.
+func _() {}}