Various performance improvements

cmd/capabilities.go

@@ -2,7 +2,6 @@ package cmd
 
 import (
 	"encoding/json"
-	"fmt"
 	"os"
 
 	"github.com/spf13/cobra"
@@ -18,13 +17,11 @@ func init() {
 		Long:   "Show capabilities for Regal",
 		RunE: func(*cobra.Command, []string) error {
 			bs, err := json.MarshalIndent(io.Capabilities(), "", "  ")
-			if err != nil {
-				return fmt.Errorf("failed marshalling capabilities: %w", err)
+			if err == nil {
+				_, err = os.Stdout.Write(append(bs, '\n'))
 			}
 
-			fmt.Fprintln(os.Stdout, string(bs))
-
-			return nil
+			return err
 		},
 	}
 

cmd/parse.go

@@ -62,12 +62,5 @@ func parse(args []string) error {
 		return err
 	}
 
-	output, err := encoding.JSON().MarshalIndent(enhancedAST, "", "  ")
-	if err != nil {
-		return err
-	}
-
-	_, err = os.Stdout.Write(output)
-
-	return err
+	return encoding.NewIndentEncoder(os.Stdout, "", "  ").Encode(enhancedAST)
 }

internal/capabilities/capabilities.go

@@ -221,9 +221,7 @@ func semverSort(stringVersions []string) {
 	}
 
 	if len(invalid) > 0 {
-		slices.Sort(invalid)
-		slices.Reverse(invalid)
-		copy(stringVersions[len(versions):], invalid)
+		copy(stringVersions[len(versions):], util.Reversed(util.Sorted(invalid)))
 	}
 }
 

internal/io/files/filter/filter.go

@@ -23,7 +23,7 @@ func Directories(_ string, info os.DirEntry) bool {
 	return info.IsDir()
 }
 
-// Filenames filters files by their exact name, not counting the path.
+// Filenames filters files by their exact name, not including the path.
 func Filenames(names ...string) Func {
 	return func(_ string, info os.DirEntry) bool {
 		return slices.ContainsFunc(names, func(name string) bool {
@@ -34,9 +34,9 @@ func Filenames(names ...string) Func {
 
 // Suffixes filters any path that has a suffix matching any of the provided suffixes.
 func Suffixes(suffixes ...string) Func {
-	return func(_ string, info os.DirEntry) bool {
+	return func(path string, _ os.DirEntry) bool {
 		return slices.ContainsFunc(suffixes, func(suffix string) bool {
-			return strings.HasSuffix(info.Name(), suffix)
+			return strings.HasSuffix(path, suffix)
 		})
 	}
 }
@@ -49,6 +49,10 @@ func Not(filter Func) Func {
 	}
 }
 
-func RegoTests(_ string, info os.DirEntry) bool {
-	return strings.HasSuffix(info.Name(), "_test.rego") && info.Name() != "todo_test.rego"
+func RegoTests(path string, info os.DirEntry) bool {
+	return strings.HasSuffix(path, "_test.rego") && info.Name() != "todo_test.rego"
+}
+
+func NotRego(path string, _ os.DirEntry) bool {
+	return !strings.HasSuffix(path, ".rego")
 }

internal/lsp/documentsymbol/documentsymbol.go

@@ -92,16 +92,17 @@ func All(contents string, module *ast.Module, builtins map[string]*ast.Builtin)
 }
 
 func locationToRange(location *ast.Location) types.Range {
-	lines := bytes.Split(location.Text, []byte("\n"))
-	numLines := len(lines)
 	startLine := util.SafeIntToUint(location.Row - 1)
+	numLines := bytes.Count(location.Text, []byte{'\n'}) + 1
 
 	endLine := startLine
-	if numLines != 1 {
+	if numLines > 1 {
 		endLine += util.SafeIntToUint(numLines - 1)
 	}
 
-	return types.RangeBetween(startLine, location.Col-1, endLine, len(lines[numLines-1]))
+	endLineContent := util.LineContents(location.Text, endLine-startLine)
+
+	return types.RangeBetween(startLine, location.Col-1, endLine, len(endLineContent))
 }
 
 func toWorkspaceSymbol(docSym types.DocumentSymbol, docURL string) types.WorkspaceSymbol {

internal/lsp/lint.go

@@ -72,12 +72,10 @@ func updateParse(ctx context.Context, opts updateParseOpts) (bool, error) {
 		return false, fmt.Errorf("failed to get file contents for uri %q", opts.FileURI)
 	}
 
-	lines := strings.Split(content, "\n")
-	numLines := len(lines)
-
 	options := rparse.ParserOptions()
 	options.RegoVersion = opts.RegoVersion
 
+	numLines := strings.Count(content, "\n") + 1
 	presentedFileName := uri.ToRelativePath(opts.FileURI, opts.WorkspaceRootURI)
 
 	module, err := rparse.ModuleWithOpts(presentedFileName, content, options)
@@ -125,6 +123,7 @@ func updateParse(ctx context.Context, opts updateParseOpts) (bool, error) {
 		}
 	}
 
+	lines := strings.Split(content, "\n")
 	diags := make([]types.Diagnostic, 0, len(astErrors))
 
 	for _, astError := range astErrors {

internal/lsp/server.go

@@ -346,7 +346,7 @@ func (l *LanguageServer) StartDiagnosticsWorker(ctx context.Context) {
 				l.sendFileDiagnostics(ctx, job.URI)
 
 				l.lintWorkspaceJobs <- lintWorkspaceJob{
-					Reason: fmt.Sprintf("file %s %s", job.URI, job.Reason),
+					Reason: "file " + job.URI + " " + job.Reason,
 					// this run is expected to used the cached aggregate state
 					// for other files.
 					// The aggregate state for this file will still be updated.
@@ -724,18 +724,13 @@ func (l *LanguageServer) StartCommandWorker(ctx context.Context) { //nolint:main
 					output := filepath.Join(l.workspacePath(), "output.json")
 
 					var f *os.File
-
-					f, err = os.OpenFile(output, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o755)
-					if err == nil {
+					if f, err = os.OpenFile(output, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o755); err == nil {
 						value := result.Value
 						if result.IsUndefined {
 							value = emptyStringAnyMap // undefined displays as an empty object
 						}
 
-						var jsonVal []byte
-						if jsonVal, err = encoding.JSON().MarshalIndent(value, "", "  "); err == nil {
-							_, err = f.Write(jsonVal)
-						}
+						err = encoding.NewIndentEncoder(f, "", "  ").Encode(value)
 
 						rio.CloseIgnore(f)
 					}

internal/roast/encoding/array.go

@@ -19,17 +19,13 @@ func (*arrayCodec) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
 
 	stream.WriteArrayStart()
 
-	i := 0
-
-	arr.Foreach(func(term *ast.Term) {
+	for i := range arr.Len() {
 		if i > 0 {
 			stream.WriteMore()
 		}
 
-		stream.WriteVal(term)
-
-		i++
-	})
+		stream.WriteVal(arr.Elem(i))
+	}
 
 	stream.WriteArrayEnd()
 }

internal/roast/encoding/head_test.go

@@ -14,37 +14,12 @@ func TestRuleHeadEncoding(t *testing.T) {
 	head := ast.Head{
 		Name: "omitted",
 		Reference: ast.Ref{
-			{
-				Value: ast.Var("foo"),
-				Location: &ast.Location{
-					Row:  1,
-					Col:  1,
-					Text: []byte("foo"),
-				},
-			},
-			{
-				Value: ast.String("bar"),
-				Location: &ast.Location{
-					Row:  1,
-					Col:  5, // following "foo."
-					Text: []byte("bar"),
-				},
-			},
-		},
-		Value: &ast.Term{
-			Value: ast.Boolean(true),
-			Location: &ast.Location{
-				Row:  1,
-				Col:  12, // following "foo.bar := "
-				Text: []byte("true"),
-			},
-		},
-		Assign: true,
-		Location: &ast.Location{
-			Row:  1,
-			Col:  1,
-			Text: []byte("foo.bar := true"),
+			{Value: ast.Var("foo"), Location: &ast.Location{Row: 1, Col: 1, Text: []byte("foo")}},
+			{Value: ast.String("bar"), Location: &ast.Location{Row: 1, Col: 5, Text: []byte("bar")}},
 		},
+		Value:    &ast.Term{Value: ast.Boolean(true), Location: &ast.Location{Row: 1, Col: 12, Text: []byte("true")}},
+		Assign:   true,
+		Location: &ast.Location{Row: 1, Col: 1, Text: []byte("foo.bar := true")},
 	}
 
 	bs, err := jsoniter.ConfigFastest.MarshalIndent(head, "", "  ")

internal/roast/encoding/location.go

@@ -1,65 +1,23 @@
 package encoding
 
 import (
-	"bytes"
-	"strconv"
-	"strings"
-	"sync"
 	"unsafe"
 
 	jsoniter "github.com/json-iterator/go"
 
 	"github.com/open-policy-agent/opa/v1/ast"
+
+	"github.com/open-policy-agent/regal/pkg/roast/rast"
 )
 
 type locationCodec struct{}
 
-var newLine = []byte("\n")
-
-var sbPool = sync.Pool{
-	New: func() any {
-		return new(strings.Builder)
-	},
-}
-
 func (*locationCodec) IsEmpty(_ unsafe.Pointer) bool {
 	return false
 }
 
 func (*locationCodec) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
-	location := *((*ast.Location)(ptr))
-
-	var endRow, endCol int
-	if location.Text == nil {
-		endRow = location.Row
-		endCol = location.Col
-	} else {
-		lines := bytes.Split(location.Text, newLine)
-
-		numLines := len(lines)
-
-		endRow = location.Row + numLines - 1
-
-		if numLines == 1 {
-			endCol = location.Col + len(location.Text)
-		} else {
-			lastLine := lines[numLines-1]
-			endCol = len(lastLine) + 1
-		}
-	}
-
-	sb := sbPool.Get().(*strings.Builder) //nolint:forcetypeassert
-
-	sb.WriteString(strconv.Itoa(location.Row))
-	sb.WriteByte(':')
-	sb.WriteString(strconv.Itoa(location.Col))
-	sb.WriteByte(':')
-	sb.WriteString(strconv.Itoa(endRow))
-	sb.WriteByte(':')
-	sb.WriteString(strconv.Itoa(endCol))
-
-	stream.WriteString(sb.String())
+	location := (*ast.Location)(ptr)
 
-	sb.Reset()
-	sbPool.Put(sb)
+	stream.SetBuffer(append(rast.AppendLocation(append(stream.Buffer(), '"'), location), '"'))
 }

internal/roast/encoding/location_test.go

@@ -16,41 +16,28 @@ func TestLocation(t *testing.T) {
 		name     string
 		location ast.Location
 		expected string
-	}{
-		{
-			name: "multiple lines",
-			location: ast.Location{
-				Row:  5,
-				Col:  2,
-				Text: []byte("allow if {\n	input.foo == true\n}"),
-			},
-			expected: "5:2:7:2",
-		},
-		{
-			name: "single line",
-			location: ast.Location{
-				Row:  1,
-				Col:  1,
-				Text: []byte("package example"),
-			},
-			expected: "1:1:1:16",
-		},
-	}
-
-	json := jsoniter.ConfigFastest
+	}{{
+		name:     "multiple lines",
+		location: ast.Location{Row: 5, Col: 2, Text: []byte("allow if {\n	input.foo == true\n}")},
+		expected: "5:2:7:2",
+	}, {
+		name:     "single line",
+		location: ast.Location{Row: 1, Col: 1, Text: []byte("package example")},
+		expected: "1:1:1:16",
+	}}
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
-			stream := json.BorrowStream(nil)
-			defer json.ReturnStream(stream)
-
+			stream := jsoniter.ConfigFastest.BorrowStream(nil)
 			stream.WriteVal(tc.location)
 
 			if string(stream.Buffer()) != fmt.Sprintf("\"%s\"", tc.expected) {
 				t.Fatalf("expected %s but got %s", tc.expected, string(stream.Buffer()))
 			}
+
+			jsoniter.ConfigFastest.ReturnStream(stream)
 		})
 	}
 }
@@ -60,10 +47,7 @@ func TestLocationHeadValue(t *testing.T) {
 	// e.g. the end column would be presented as before the start column.
 	t.Parallel()
 
-	module := ast.MustParseModule("package foo.bar\n\nrule := true")
-	json := jsoniter.ConfigFastest
-
-	out, err := json.MarshalIndent(module, "", "  ")
+	out, err := jsoniter.ConfigFastest.MarshalIndent(ast.MustParseModule("package foo.bar\n\nrule := true"), "", "  ")
 	if err != nil {
 		t.Fatalf("failed to marshal module: %v", err)
 	}

internal/roast/encoding/module_test.go

@@ -228,8 +228,7 @@ func TestSerializedModuleSize(t *testing.T) {
 	}
 }
 
-// BenchmarkSerializeModule-16    	    3361	    354640 ns/op	  216756 B/op	    9773 allocs/op
-
+// 285329 ns/op	  125555 B/op	    3094 allocs/op
 func BenchmarkSerializeModule(b *testing.B) {
 	policy := mustReadTestFile(b, "testdata/policy.rego")
 	module := ast.MustParseModuleWithOpts(string(policy), ast.ParserOptions{ProcessAnnotation: true})