internal/lsp: add foldingRange support

Support textDocument/foldingRange request. Provide folding ranges for
multiline comment blocks, declarations, block statements, field lists,
case clauses, and call expressions.

Fixes golang/go#32987

Change-Id: I9c76e850ffa0e5bb65bee273d8ee40577c342f92
Reviewed-on: https://go-review.googlesource.com/c/tools/+/192257
Run-TryBot: Suzy Mueller <suzmue@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

internal/lsp/cmd/cmd_test.go

@@ -44,6 +44,10 @@ func (r *runner) Completion(t *testing.T, data tests.Completions, snippets tests
 	//TODO: add command line completions tests when it works
 }
 
+func (r *runner) FoldingRange(t *testing.T, data tests.FoldingRanges) {
+	//TODO: add command line folding range tests when it works
+}
+
 func (r *runner) Highlight(t *testing.T, data tests.Highlights) {
 	//TODO: add command line highlight tests when it works
 }

internal/lsp/folding_range.go

@@ -0,0 +1,28 @@
+package lsp
+
+import (
+	"context"
+
+	"golang.org/x/tools/internal/lsp/protocol"
+	"golang.org/x/tools/internal/lsp/source"
+	"golang.org/x/tools/internal/span"
+)
+
+func (s *Server) foldingRange(ctx context.Context, params *protocol.FoldingRangeParams) ([]protocol.FoldingRange, error) {
+	uri := span.NewURI(params.TextDocument.URI)
+	view := s.session.ViewOf(uri)
+	f, err := getGoFile(ctx, view, uri)
+	if err != nil {
+		return nil, err
+	}
+	m, err := getMapper(ctx, f)
+	if err != nil {
+		return nil, err
+	}
+
+	ranges, err := source.FoldingRange(ctx, view, f)
+	if err != nil {
+		return nil, err
+	}
+	return source.ToProtocolFoldingRanges(m, ranges)
+}

internal/lsp/general.go

@@ -99,6 +99,7 @@ func (s *Server) initialize(ctx context.Context, params *protocol.InitializePara
 			DefinitionProvider:         true,
 			DocumentFormattingProvider: true,
 			DocumentSymbolProvider:     true,
+			FoldingRangeProvider:       true,
 			HoverProvider:              true,
 			DocumentHighlightProvider:  true,
 			DocumentLinkProvider:       &protocol.DocumentLinkOptions{},

internal/lsp/lsp_test.go

@@ -254,6 +254,103 @@ func summarizeCompletionItems(i int, want []source.CompletionItem, got []protoco
 	return msg.String()
 }
 
+func (r *runner) FoldingRange(t *testing.T, data tests.FoldingRanges) {
+	for _, spn := range data {
+		uri := spn.URI()
+		filename := uri.Filename()
+
+		ranges, err := r.server.FoldingRange(r.ctx, &protocol.FoldingRangeParams{
+			TextDocument: protocol.TextDocumentIdentifier{
+				URI: protocol.NewURI(uri),
+			},
+		})
+		if err != nil {
+			t.Error(err)
+			continue
+		}
+
+		f, err := getGoFile(r.ctx, r.server.session.ViewOf(uri), uri)
+		if err != nil {
+			t.Fatal(err)
+		}
+		m, err := getMapper(r.ctx, f)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Fold all ranges.
+		got, err := foldRanges(m, string(m.Content), ranges)
+		if err != nil {
+			t.Error(err)
+			continue
+		}
+		want := string(r.data.Golden("foldingRange", spn.URI().Filename(), func() ([]byte, error) {
+			return []byte(got), nil
+		}))
+
+		if want != got {
+			t.Errorf("foldingRanges failed for %s, expected:\n%v\ngot:\n%v", filename, want, got)
+		}
+
+		// Filter by kind.
+		kinds := []protocol.FoldingRangeKind{protocol.Imports, protocol.Comment}
+		for _, kind := range kinds {
+			var kindOnly []protocol.FoldingRange
+			for _, fRng := range ranges {
+				if fRng.Kind == string(kind) {
+					kindOnly = append(kindOnly, fRng)
+				}
+			}
+
+			got, err := foldRanges(m, string(m.Content), kindOnly)
+			if err != nil {
+				t.Error(err)
+				continue
+			}
+			want := string(r.data.Golden("foldingRange-"+string(kind), spn.URI().Filename(), func() ([]byte, error) {
+				return []byte(got), nil
+			}))
+
+			if want != got {
+				t.Errorf("foldingRanges-%s failed for %s, expected:\n%v\ngot:\n%v", string(kind), filename, want, got)
+			}
+
+		}
+
+	}
+}
+
+func foldRanges(m *protocol.ColumnMapper, contents string, ranges []protocol.FoldingRange) (string, error) {
+	// TODO(suzmue): Allow folding ranges to intersect for these tests, do a folding by level,
+	// or per individual fold.
+	foldedText := "<>"
+	res := contents
+	// Apply the edits from the end of the file forward
+	// to preserve the offsets
+	for i := len(ranges) - 1; i >= 0; i-- {
+		fRange := ranges[i]
+		spn, err := m.RangeSpan(protocol.Range{
+			Start: protocol.Position{
+				Line:      fRange.StartLine,
+				Character: fRange.StartCharacter,
+			},
+			End: protocol.Position{
+				Line:      fRange.EndLine,
+				Character: fRange.EndCharacter,
+			},
+		})
+		if err != nil {
+			return "", err
+		}
+		start := spn.Start().Offset()
+		end := spn.End().Offset()
+
+		tmp := res[0:start] + foldedText
+		res = tmp + res[end:]
+	}
+	return res, nil
+}
+
 func (r *runner) Format(t *testing.T, data tests.Formats) {
 	for _, spn := range data {
 		uri := spn.URI()

internal/lsp/server.go

@@ -260,8 +260,8 @@ func (s *Server) Declaration(context.Context, *protocol.TextDocumentPositionPara
 	return nil, notImplemented("Declaration")
 }
 
-func (s *Server) FoldingRange(context.Context, *protocol.FoldingRangeParams) ([]protocol.FoldingRange, error) {
-	return nil, notImplemented("FoldingRange")
+func (s *Server) FoldingRange(ctx context.Context, params *protocol.FoldingRangeParams) ([]protocol.FoldingRange, error) {
+	return s.foldingRange(ctx, params)
 }
 
 func (s *Server) LogTraceNotification(context.Context, *protocol.LogTraceParams) error {

internal/lsp/source/folding_range.go

@@ -0,0 +1,118 @@
+package source
+
+import (
+	"context"
+	"go/ast"
+	"go/token"
+	"sort"
+
+	"golang.org/x/tools/internal/lsp/protocol"
+	"golang.org/x/tools/internal/span"
+)
+
+type FoldingRangeInfo struct {
+	Range span.Range
+	Kind  protocol.FoldingRangeKind
+}
+
+// FoldingRange gets all of the folding range for f.
+func FoldingRange(ctx context.Context, view View, f GoFile) (ranges []FoldingRangeInfo, err error) {
+	// TODO(suzmue): consider limiting the number of folding ranges returned, and
+	// implement a way to prioritize folding ranges in that case.
+	file, err := f.GetAST(ctx, ParseFull)
+	if err != nil {
+		return nil, err
+	}
+
+	// Get folding ranges for comments separately as they are not walked by ast.Inspect.
+	ranges = append(ranges, commentsFoldingRange(f.FileSet(), file)...)
+
+	visit := func(n ast.Node) bool {
+		var kind protocol.FoldingRangeKind
+		var start, end token.Pos
+		switch n := n.(type) {
+		case *ast.BlockStmt:
+			// Fold from position of "{" to position of "}".
+			start, end = n.Lbrace+1, n.Rbrace
+		case *ast.CaseClause:
+			// 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
+		case *ast.FieldList:
+			// Fold from position of opening parenthesis/brace, to position of
+			// closing parenthesis/brace.
+			start, end = n.Opening+1, n.Closing
+		case *ast.GenDecl:
+			// If this is an import declaration, set the kind to be protocol.Imports.
+			if n.Tok == token.IMPORT {
+				kind = protocol.Imports
+			}
+			// Fold from position of "(" to position of ")".
+			start, end = n.Lparen+1, n.Rparen
+		}
+
+		if start.IsValid() && end.IsValid() {
+			ranges = append(ranges, FoldingRangeInfo{
+				Range: span.NewRange(f.FileSet(), start, end),
+				Kind:  kind,
+			})
+		}
+		return true
+	}
+
+	// Walk the ast and collect folding ranges.
+	ast.Inspect(file, visit)
+
+	sort.Slice(ranges, func(i, j int) bool {
+		if ranges[i].Range.Start < ranges[j].Range.Start {
+			return true
+		} else if ranges[i].Range.Start > ranges[j].Range.Start {
+			return false
+		}
+		return ranges[i].Range.End < ranges[j].Range.End
+	})
+	return ranges, nil
+}
+
+// commentsFoldingRange returns the folding ranges for all comment blocks in file.
+// The folding range starts at the end of the first comment, and ends at the end of the
+// comment block and has kind protocol.Comment.
+func commentsFoldingRange(fset *token.FileSet, file *ast.File) []FoldingRangeInfo {
+	var comments []FoldingRangeInfo
+	for _, commentGrp := range file.Comments {
+		// Don't fold single comments.
+		if len(commentGrp.List) <= 1 {
+			continue
+		}
+		comments = append(comments, FoldingRangeInfo{
+			// Fold from the end of the first line comment to the end of the comment block.
+			Range: span.NewRange(fset, commentGrp.List[0].End(), commentGrp.End()),
+			Kind:  protocol.Comment,
+		})
+	}
+	return comments
+}
+
+func ToProtocolFoldingRanges(m *protocol.ColumnMapper, ranges []FoldingRangeInfo) ([]protocol.FoldingRange, error) {
+	var res []protocol.FoldingRange
+	for _, r := range ranges {
+		spn, err := r.Range.Span()
+		if err != nil {
+			return nil, err
+		}
+		rng, err := m.Range(spn)
+		if err != nil {
+			return nil, err
+		}
+		res = append(res, protocol.FoldingRange{
+			StartLine:      rng.Start.Line,
+			StartCharacter: rng.Start.Character,
+			EndLine:        rng.End.Line,
+			EndCharacter:   rng.End.Character,
+			Kind:           string(r.Kind),
+		})
+	}
+	return res, nil
+}

internal/lsp/source/source_test.go

@@ -257,6 +257,89 @@ func summarizeCompletionItems(i int, want []source.CompletionItem, got []source.
 	return msg.String()
 }
 
+func (r *runner) FoldingRange(t *testing.T, data tests.FoldingRanges) {
+	for _, spn := range data {
+		uri := spn.URI()
+		filename := uri.Filename()
+
+		f, err := r.view.GetFile(r.ctx, uri)
+		if err != nil {
+			t.Fatalf("failed for %v: %v", spn, err)
+		}
+
+		ranges, err := source.FoldingRange(r.ctx, r.view, f.(source.GoFile))
+		if err != nil {
+			t.Error(err)
+			continue
+		}
+		data, _, err := f.Handle(r.ctx).Read(r.ctx)
+		if err != nil {
+			t.Error(err)
+			continue
+		}
+		// Fold all ranges.
+		got, err := foldRanges(string(data), ranges)
+		if err != nil {
+			t.Error(err)
+			continue
+		}
+		want := string(r.data.Golden("foldingRange", spn.URI().Filename(), func() ([]byte, error) {
+			return []byte(got), nil
+		}))
+
+		if want != got {
+			t.Errorf("foldingRanges failed for %s, expected:\n%v\ngot:\n%v", filename, want, got)
+		}
+
+		// Filter by kind.
+		kinds := []protocol.FoldingRangeKind{protocol.Imports, protocol.Comment}
+		for _, kind := range kinds {
+			var kindOnly []source.FoldingRangeInfo
+			for _, fRng := range ranges {
+				if fRng.Kind == kind {
+					kindOnly = append(kindOnly, fRng)
+				}
+			}
+
+			got, err := foldRanges(string(data), kindOnly)
+			if err != nil {
+				t.Error(err)
+				continue
+			}
+			want := string(r.data.Golden("foldingRange-"+string(kind), spn.URI().Filename(), func() ([]byte, error) {
+				return []byte(got), nil
+			}))
+
+			if want != got {
+				t.Errorf("foldingRanges-%s failed for %s, expected:\n%v\ngot:\n%v", string(kind), filename, want, got)
+			}
+
+		}
+
+	}
+}
+
+func foldRanges(contents string, ranges []source.FoldingRangeInfo) (string, error) {
+	// TODO(suzmue): Allow folding ranges to intersect for these tests.
+	foldedText := "<>"
+	res := contents
+	// Apply the folds from the end of the file forward
+	// to preserve the offsets.
+	for i := len(ranges) - 1; i >= 0; i-- {
+		fRange := ranges[i]
+		spn, err := fRange.Range.Span()
+		if err != nil {
+			return "", err
+		}
+		start := spn.Start().Offset()
+		end := spn.End().Offset()
+
+		tmp := res[0:start] + foldedText
+		res = tmp + res[end:]
+	}
+	return res, nil
+}
+
 func (r *runner) Format(t *testing.T, data tests.Formats) {
 	ctx := r.ctx
 	for _, spn := range data {

internal/lsp/testdata/folding/a.go

@@ -0,0 +1,17 @@
+package folding //@fold("package")
+
+import (
+	_ "fmt"
+	_ "log"
+)
+
+import _ "os"
+
+// bar is a function.
+// With a multiline doc comment.
+func bar() string {
+	return `
+this string
+is not indented`
+
+}

internal/lsp/testdata/folding/a.go.golden

@@ -0,0 +1,44 @@
+-- foldingRange --
+package folding //@fold("package")
+
+import (<>)
+
+import _ "os"
+
+// bar is a function.<>
+func bar(<>) string {<>}
+
+-- foldingRange-comment --
+package folding //@fold("package")
+
+import (
+	_ "fmt"
+	_ "log"
+)
+
+import _ "os"
+
+// bar is a function.<>
+func bar() string {
+	return `
+this string
+is not indented`
+
+}
+
+-- foldingRange-imports --
+package folding //@fold("package")
+
+import (<>)
+
+import _ "os"
+
+// bar is a function.
+// With a multiline doc comment.
+func bar() string {
+	return `
+this string
+is not indented`
+
+}
+

internal/lsp/tests/tests.go

@@ -36,6 +36,7 @@ const (
 	ExpectedImportCount            = 2
 	ExpectedDefinitionsCount       = 39
 	ExpectedTypeDefinitionsCount   = 2
+	ExpectedFoldingRangesCount     = 1
 	ExpectedHighlightsCount        = 2
 	ExpectedReferencesCount        = 5
 	ExpectedRenamesCount           = 20
@@ -57,6 +58,7 @@ type Diagnostics map[span.URI][]source.Diagnostic
 type CompletionItems map[token.Pos]*source.CompletionItem
 type Completions map[span.Span][]token.Pos
 type CompletionSnippets map[span.Span]CompletionSnippet
+type FoldingRanges []span.Span
 type Formats []span.Span
 type Imports []span.Span
 type Definitions map[span.Span]Definition
@@ -75,6 +77,7 @@ type Data struct {
 	CompletionItems    CompletionItems
 	Completions        Completions
 	CompletionSnippets CompletionSnippets
+	FoldingRanges      FoldingRanges
 	Formats            Formats
 	Imports            Imports
 	Definitions        Definitions
@@ -95,6 +98,7 @@ type Data struct {
 type Tests interface {
 	Diagnostics(*testing.T, Diagnostics)
 	Completion(*testing.T, Completions, CompletionSnippets, CompletionItems)
+	FoldingRange(*testing.T, FoldingRanges)
 	Format(*testing.T, Formats)
 	Import(*testing.T, Imports)
 	Definition(*testing.T, Definitions)
@@ -222,6 +226,7 @@ func Load(t testing.TB, exporter packagestest.Exporter, dir string) *Data {
 		"diag":      data.collectDiagnostics,
 		"item":      data.collectCompletionItems,
 		"complete":  data.collectCompletions,
+		"fold":      data.collectFoldingRanges,
 		"format":    data.collectFormats,
 		"import":    data.collectImports,
 		"godef":     data.collectDefinitions,
@@ -278,6 +283,14 @@ func Run(t *testing.T, tests Tests, data *Data) {
 		tests.Diagnostics(t, data.Diagnostics)
 	})
 
+	t.Run("FoldingRange", func(t *testing.T) {
+		t.Helper()
+		if len(data.FoldingRanges) != ExpectedFoldingRangesCount {
+			t.Errorf("got %v folding ranges expected %v", len(data.FoldingRanges), ExpectedFoldingRangesCount)
+		}
+		tests.FoldingRange(t, data.FoldingRanges)
+	})
+
 	t.Run("Format", func(t *testing.T) {
 		t.Helper()
 		if len(data.Formats) != ExpectedFormatCount {
@@ -547,6 +560,10 @@ func (data *Data) collectCompletionItems(pos token.Pos, args []string) {
 	}
 }
 
+func (data *Data) collectFoldingRanges(spn span.Span) {
+	data.FoldingRanges = append(data.FoldingRanges, spn)
+}
+
 func (data *Data) collectFormats(spn span.Span) {
 	data.Formats = append(data.Formats, spn)
 }