lsp: Support folding ranges (#663)

Adding folding range support for:
* Imports
* Comments (when grouped on more than one line)
* Blocks from tokens like `{}`, `[]`, `()`

Fixes #648

Signed-off-by: Anders Eknert <anders@styra.com>

.github/workflows/build.yaml

@@ -42,7 +42,7 @@ jobs:
       - uses: golangci/golangci-lint-action@v4.0.0
         if: matrix.os.name == 'linux'
         with:
-          version: v1.56.1
+          version: v1.57.2
       - uses: actions/upload-artifact@v4
         with:
           name: regal-${{ matrix.os.name }}

.golangci.yaml

@@ -61,3 +61,7 @@ linters-settings:
       - prefix(github.com/styrainc/regal)
       - blank
       - dot
+
+issues:
+  exclude-dirs:
+    - internal/lsp/opa

e2e/cli_test.go

@@ -746,7 +746,7 @@ allow if {
 	true
 }
 `)
-	err := os.WriteFile(filepath.Join(td, "main.rego"), unformattedContents, 0644)
+	err := os.WriteFile(filepath.Join(td, "main.rego"), unformattedContents, 0o644)
 	if err != nil {
 		t.Fatalf("failed to write main.rego: %v", err)
 	}

internal/lsp/foldingrange.go

@@ -0,0 +1,155 @@
+package lsp
+
+import (
+	"strings"
+
+	"github.com/open-policy-agent/opa/ast"
+
+	"github.com/styrainc/regal/internal/lsp/opa/scanner"
+	"github.com/styrainc/regal/internal/lsp/opa/tokens"
+	"github.com/styrainc/regal/internal/lsp/types"
+)
+
+type stack []scanner.Position
+
+func (s stack) Push(p scanner.Position) stack {
+	return append(s, p)
+}
+
+func (s stack) Pop() (stack, scanner.Position) {
+	l := len(s)
+
+	return s[:l-1], s[l-1]
+}
+
+func TokenFoldingRanges(policy string) []types.FoldingRange {
+	scn, err := scanner.New(strings.NewReader(policy))
+	if err != nil {
+		panic(err)
+	}
+
+	var lastPosition scanner.Position
+
+	curlyBraceStack := stack{}
+	bracketStack := stack{}
+	parensStack := stack{}
+
+	foldingRanges := make([]types.FoldingRange, 0)
+
+	for {
+		token, position, _, errors := scn.Scan()
+
+		if token == tokens.EOF || len(errors) > 0 {
+			break
+		}
+
+		switch {
+		case token == tokens.LBrace:
+			curlyBraceStack = curlyBraceStack.Push(position)
+		case token == tokens.RBrace && len(curlyBraceStack) > 0:
+			curlyBraceStack, lastPosition = curlyBraceStack.Pop()
+
+			startChar := uint(lastPosition.Col)
+
+			foldingRanges = append(foldingRanges, types.FoldingRange{
+				StartLine:      uint(lastPosition.Row - 1),
+				StartCharacter: &startChar,
+				// Note that we stop at the line _before_ the closing curly brace
+				// as that shows the end of the object/set in the client, which seems
+				// to be how other implementations do it
+				EndLine: uint(position.Row - 2),
+				Kind:    "region",
+			})
+		case token == tokens.LBrack:
+			bracketStack = bracketStack.Push(position)
+		case token == tokens.RBrack && len(bracketStack) > 0:
+			bracketStack, lastPosition = bracketStack.Pop()
+
+			startChar := uint(lastPosition.Col)
+
+			foldingRanges = append(foldingRanges, types.FoldingRange{
+				StartLine:      uint(lastPosition.Row - 1),
+				StartCharacter: &startChar,
+				// Note that we stop at the line _before_ the closing bracket
+				// as that shows the end of the array in the client, which seems
+				// to be how other implementations do it
+				EndLine: uint(position.Row - 2),
+				Kind:    "region",
+			})
+		case token == tokens.LParen:
+			parensStack = parensStack.Push(position)
+		case token == tokens.RParen && len(parensStack) > 0:
+			parensStack, lastPosition = parensStack.Pop()
+
+			startChar := uint(lastPosition.Col)
+
+			foldingRanges = append(foldingRanges, types.FoldingRange{
+				StartLine:      uint(lastPosition.Row - 1),
+				StartCharacter: &startChar,
+				// Note that we stop at the line _before_ the closing bracket
+				// as that shows the end of the array in the client, which seems
+				// to be how other implementations do it
+				EndLine: uint(position.Row - 2),
+				Kind:    "region",
+			})
+		}
+	}
+
+	return foldingRanges
+}
+
+func findFoldingRanges(text string, module *ast.Module) []types.FoldingRange {
+	uintZero := uint(0)
+
+	ranges := make([]types.FoldingRange, 0)
+
+	// Comments
+
+	numComments := len(module.Comments)
+	isBlock := false
+
+	var startLine uint
+
+	for i, comment := range module.Comments {
+		// the comment following this is on the next line
+		if i+1 < numComments && module.Comments[i+1].Location.Row == comment.Location.Row+1 {
+			isBlock = true
+
+			if i == 0 || module.Comments[i-1].Location.Row != comment.Location.Row-1 {
+				startLine = uint(comment.Location.Row - 1)
+			}
+		} else if isBlock {
+			ranges = append(ranges, types.FoldingRange{
+				StartLine:      startLine,
+				EndLine:        uint(comment.Location.Row - 1),
+				StartCharacter: &uintZero,
+				Kind:           "comment",
+			})
+
+			isBlock = false
+		}
+	}
+
+	// Imports
+
+	if len(module.Imports) > 2 {
+		lastImport := module.Imports[len(module.Imports)-1]
+
+		// note that we treat *all* imports as a single folding range,
+		// as it's likely the user wants to hide all of them... but we
+		// could consider instead folding "blocks" of grouped imports
+
+		ranges = append(ranges, types.FoldingRange{
+			StartLine:      uint(module.Imports[0].Location.Row - 1),
+			EndLine:        uint(lastImport.Location.Row - 1),
+			StartCharacter: &uintZero,
+			Kind:           "imports",
+		})
+	}
+
+	// Tokens — {}, [], ()
+
+	tokenRanges := TokenFoldingRanges(text)
+
+	return append(ranges, tokenRanges...)
+}

internal/lsp/foldingrange_test.go

@@ -0,0 +1,76 @@
+package lsp
+
+import (
+	"testing"
+)
+
+func TestTokenFoldingRanges(t *testing.T) {
+	t.Parallel()
+
+	policy := `package p
+
+	import rego.v1
+
+rule if {
+	arr := [
+		1,
+		2,
+		3,
+	]
+	par := (
+		1 +
+		2 -
+		3
+	)
+}`
+
+	foldingRanges := TokenFoldingRanges(policy)
+
+	if len(foldingRanges) != 3 {
+		t.Fatalf("Expected 3 folding ranges, got %d", len(foldingRanges))
+	}
+
+	arr := foldingRanges[0]
+
+	if arr.StartLine != 5 || *arr.StartCharacter != 9 {
+		t.Errorf("Expected start line 5 and start character 9, got %d and %d", arr.StartLine, *arr.StartCharacter)
+	}
+
+	if arr.EndLine != 8 {
+		t.Errorf("Expected end line 8, got %d", arr.EndLine)
+	}
+
+	parens := foldingRanges[1]
+
+	if parens.StartLine != 10 || *parens.StartCharacter != 9 {
+		t.Errorf("Expected start line 10 and start character 9, got %d and %d", parens.StartLine, *parens.StartCharacter)
+	}
+
+	if parens.EndLine != 13 {
+		t.Errorf("Expected end line 13, got %d", parens.EndLine)
+	}
+
+	rule := foldingRanges[2]
+
+	if rule.StartLine != 4 || *rule.StartCharacter != 9 {
+		t.Errorf("Expected start line 4 and start character 9, got %d and %d", rule.StartLine, *rule.StartCharacter)
+	}
+
+	if rule.EndLine != 14 {
+		t.Errorf("Expected end line 7, got %d", rule.EndLine)
+	}
+}
+
+func TestTokenInvalidFoldingRanges(t *testing.T) {
+	t.Parallel()
+
+	policy := `package p
+
+arr := ]]`
+
+	foldingRanges := TokenFoldingRanges(policy)
+
+	if len(foldingRanges) != 0 {
+		t.Fatalf("Expected no folding ranges, got %d", len(foldingRanges))
+	}
+}