perf: avoid JSON/map roundtrips in ast.Value transforms

bundle/regal/lsp/clients/clients.rego

@@ -1,3 +1,5 @@
+# METADATA
+# description: Client identifiers known by the Regal LSP server
 package regal.lsp.clients
 
 # METADATA

bundle/regal/lsp/completion/providers/packagename/packagename.rego

@@ -16,7 +16,7 @@ items contains item if {
 	startswith(line, "package ")
 	position.character > 7
 
-	ps := input.regal.context.path_separator
+	ps := input.regal.environment.path_separator
 
 	abs_dir := _base(input.regal.file.name)
 	rel_dir := trim_prefix(abs_dir, input.regal.context.workspace_root)

bundle/regal/lsp/completion/providers/packagename/packagename_test.rego

@@ -10,13 +10,13 @@ test_package_name_completion_on_typing if {
 			"lines": split(policy, "\n"),
 		},
 		"context": {
-			"path_separator": "/",
 			"workspace_root": "/Users/joe/policy",
 			"location": {
 				"row": 1,
 				"col": 10,
 			},
 		},
+		"environment": {"path_separator": "/"},
 	}}
 	items := provider.items with input as provider_input
 	items == {{
@@ -41,13 +41,13 @@ test_package_name_completion_on_typing_multiple_suggestions if {
 			"lines": split(policy, "\n"),
 		},
 		"context": {
-			"path_separator": "/",
 			"workspace_root": "/Users/joe/policy",
 			"location": {
 				"row": 1,
 				"col": 10,
 			},
 		},
+		"environment": {"path_separator": "/"},
 	}}
 	items := provider.items with input as provider_input
 	items == {
@@ -86,13 +86,13 @@ test_package_name_completion_on_typing_multiple_suggestions_when_invoked if {
 			"lines": split(policy, "\n"),
 		},
 		"context": {
-			"path_separator": "/",
 			"workspace_root": "/Users/joe/policy",
 			"location": {
 				"row": 1,
 				"col": 9,
 			},
 		},
+		"environment": {"path_separator": "/"},
 	}}
 	items := provider.items with input as provider_input
 	items == {

bundle/regal/lsp/completion/providers/regov1/regov1.rego

@@ -8,9 +8,9 @@ import data.regal.lsp.completion.kind
 import data.regal.lsp.completion.location
 
 # METADATA
-# description: completion suggestion for rego.v1
+# description: completion suggestion for rego.v1 import
 items contains item if {
-	input.regal.context.rego_version != 3 # the rego.v1 import is not used in v1 rego (3)
+	input.regal.file.rego_version != "v1" # the rego.v1 import is not used in v1 Rego
 	not strings.any_prefix_match(input.regal.file.lines, "import rego.v1")
 
 	position := location.to_position(input.regal.context.location)

bundle/regal/lsp/completion/providers/regov1/regov1_test.rego

@@ -8,7 +8,7 @@ test_regov1_completion_on_typing if {
 	policy := `package policy
 
 import r`
-	items := provider.items with input as util.input_with_location_and_version(policy, {"row": 3, "col": 9}, 0)
+	items := provider.items with input as util.input_with_location_and_version(policy, {"row": 3, "col": 9}, "v0")
 	items == {{
 		"label": "rego.v1",
 		"kind": 9,
@@ -33,7 +33,7 @@ test_regov1_completion_on_invoked if {
 	policy := `package policy
 
 import `
-	items := provider.items with input as util.input_with_location_and_version(policy, {"row": 3, "col": 8}, 0)
+	items := provider.items with input as util.input_with_location_and_version(policy, {"row": 3, "col": 8}, "v0")
 	items == {{
 		"label": "rego.v1",
 		"kind": 9,
@@ -60,7 +60,7 @@ test_no_regov1_completion_if_already_imported if {
 import rego.v1
 
 import r`
-	items := provider.items with input as util.input_with_location_and_version(policy, {"row": 5, "col": 9}, 0)
+	items := provider.items with input as util.input_with_location_and_version(policy, {"row": 5, "col": 9}, "v0")
 	items == set()
 }
 
@@ -71,7 +71,7 @@ import r`
 	items := provider.items with input as util.input_with_location_and_version(
 		policy,
 		{"row": 3, "col": 9},
-		3, # RegoV1
+		"v1", # RegoV1
 	)
 	items == set()
 }

bundle/regal/lsp/completion/providers/test_utils/test_utils.rego

@@ -35,9 +35,7 @@ input_with_location_and_version(policy, location, rego_version) := {"regal": {
 	"file": {
 		"name": "p.rego",
 		"lines": split(policy, "\n"),
-	},
-	"context": {
-		"location": location,
 		"rego_version": rego_version,
 	},
+	"context": {"location": location},
 }}

cmd/test.go

@@ -397,7 +397,7 @@ func Runtime() *ast.Term {
 	for _, s := range os.Environ() {
 		parts := strings.SplitN(s, "=", 2)
 		if len(parts) == 1 {
-			env.Insert(ast.StringTerm(parts[0]), ast.NullTerm())
+			env.Insert(ast.StringTerm(parts[0]), ast.InternedNullTerm)
 		} else if len(parts) > 1 {
 			env.Insert(ast.StringTerm(parts[0]), ast.StringTerm(parts[1]))
 		}

go.mod

@@ -22,7 +22,7 @@ require (
 	github.com/sourcegraph/jsonrpc2 v0.2.1
 	github.com/spf13/cobra v1.9.1
 	github.com/spf13/pflag v1.0.6
-	github.com/styrainc/roast v0.12.0
+	github.com/styrainc/roast v0.14.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 

go.sum

@@ -305,8 +305,8 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
 github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-github.com/styrainc/roast v0.12.0 h1:TCxazAd1PodsmzMlWgaomDxsx48QdFqJWGldxC7YEpg=
-github.com/styrainc/roast v0.12.0/go.mod h1:cKELz96vKP1jeA/0HeyMW6gOqvDI6eo3wmq5/a43TYA=
+github.com/styrainc/roast v0.14.0 h1:UewC6U8A2MvUUnXRQggHFvpr94jSHzFEvcyYVfbLutc=
+github.com/styrainc/roast v0.14.0/go.mod h1:cKELz96vKP1jeA/0HeyMW6gOqvDI6eo3wmq5/a43TYA=
 github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8=
 github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
 github.com/tchap/go-patricia/v2 v2.3.2 h1:xTHFutuitO2zqKAQ5rCROYgUb7Or/+IC3fts9/Yc7nM=

internal/lsp/clients/clients.go

@@ -2,7 +2,7 @@ package clients
 
 // Identifier represent different supported clients and can be used to toggle or change
 // server behavior based on the client.
-type Identifier int
+type Identifier uint8
 
 const (
 	IdentifierGeneric Identifier = iota

internal/lsp/completions/providers/policy.go

@@ -18,6 +18,7 @@ import (
 	"github.com/styrainc/regal/pkg/builtins"
 
 	"github.com/styrainc/roast/pkg/encoding"
+	"github.com/styrainc/roast/pkg/transform"
 )
 
 // Policy provides suggestions that have been determined by Rego policy.
@@ -58,40 +59,42 @@ func (p *Policy) Run(
 		return nil, fmt.Errorf("could not get file contents for: %s", params.TextDocument.URI)
 	}
 
+	// input.regal.context
 	location := rego2.LocationFromPosition(params.Position)
-	inputContext := make(map[string]any)
-	inputContext["location"] = map[string]any{
-		"row": location.Row,
-		"col": location.Col,
-	}
-	inputContext["client_identifier"] = opts.ClientIdentifier
-	inputContext["workspace_root"] = uri.ToPath(opts.ClientIdentifier, opts.RootURI)
-	inputContext["path_separator"] = rio.PathSeparator
-	inputContext["rego_version"] = opts.RegoVersion
-
-	workspacePath := uri.ToPath(opts.ClientIdentifier, opts.RootURI)
-
-	inputDotJSONPath, inputDotJSONContent := rio.FindInput(
-		uri.ToPath(opts.ClientIdentifier, params.TextDocument.URI),
-		workspacePath,
+	regalContext := ast.NewObject(
+		ast.Item(ast.InternedStringTerm("location"), ast.ObjectTerm(
+			ast.Item(ast.InternedStringTerm("row"), ast.InternedIntNumberTerm(location.Row)),
+			ast.Item(ast.InternedStringTerm("col"), ast.InternedIntNumberTerm(location.Col)),
+		)),
+		ast.Item(ast.InternedStringTerm("client_identifier"), ast.InternedIntNumberTerm(int(opts.ClientIdentifier))),
+		ast.Item(ast.InternedStringTerm("workspace_root"), ast.InternedStringTerm(opts.RootURI)),
 	)
 
+	path := uri.ToPath(opts.ClientIdentifier, params.TextDocument.URI)
+
+	// TODO: Avoid the intermediate map[string]any step and unmarshal directly into ast.Value.
+	inputDotJSONPath, inputDotJSONContent := rio.FindInput(path, uri.ToPath(opts.ClientIdentifier, opts.RootURI))
 	if inputDotJSONPath != "" && inputDotJSONContent != nil {
-		inputContext["input_dot_json_path"] = inputDotJSONPath
-		inputContext["input_dot_json"] = inputDotJSONContent
-	}
+		inputDotJSONValue, err := transform.ToOPAInputValue(inputDotJSONContent)
+		if err != nil {
+			return nil, fmt.Errorf("failed converting input dot JSON content to value: %w", err)
+		}
 
-	input, err := rego2.ToInput(
-		params.TextDocument.URI,
-		opts.ClientIdentifier,
-		content,
-		inputContext,
-	)
-	if err != nil {
-		// parser error could be due to work in progress, so just return an empty list here
-		return []types.CompletionItem{}, nil //nolint: nilerr
+		regalContext.Insert(ast.InternedStringTerm("input_dot_json_path"), ast.InternedStringTerm(inputDotJSONPath))
+		regalContext.Insert(ast.InternedStringTerm("input_dot_json"), ast.NewTerm(inputDotJSONValue))
 	}
 
+	// input.regal
+	regalObj := transform.RegalContext(path, content, opts.RegoVersion.String())
+	regalObj.Insert(ast.InternedStringTerm("context"), ast.NewTerm(regalContext))
+
+	fileRef := ast.Ref{ast.InternedStringTerm("file")}
+	fileObj, _ := regalObj.Find(fileRef)
+	//nolint:forcetypeassert
+	fileObj.(ast.Object).Insert(ast.InternedStringTerm("uri"), ast.InternedStringTerm(params.TextDocument.URI))
+
+	input := ast.NewObject(ast.Item(ast.InternedStringTerm("regal"), ast.NewTerm(regalObj)))
+
 	result, err := rego2.QueryRegalBundle(ctx, input, p.pq)
 	if err != nil {
 		return nil, fmt.Errorf("failed querying regal bundle: %w", err)

internal/lsp/completions/providers/policy_test.go

@@ -44,28 +44,18 @@ allow if {
 	}, inmem.OptRoundTripOnWrite(false))
 
 	locals := NewPolicy(t.Context(), store)
-
 	params := types.CompletionParams{
-		TextDocument: types.TextDocumentIdentifier{
-			URI: testCaseFileURI,
-		},
-		Position: types.Position{
-			Line:      5,
-			Character: 11,
-		},
+		TextDocument: types.TextDocumentIdentifier{URI: testCaseFileURI},
+		Position:     types.Position{Line: 5, Character: 11},
 	}
+	opts := &Options{ClientIdentifier: clients.IdentifierGeneric}
 
-	result, err := locals.Run(
-		t.Context(),
-		c,
-		params,
-		&Options{ClientIdentifier: clients.IdentifierGeneric},
-	)
+	result, err := locals.Run(t.Context(), c, params, opts)
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
-	labels := []string{}
+	labels := make([]string, 0, len(result))
 	for _, item := range result {
 		labels = append(labels, item.Label)
 	}

internal/lsp/completions/providers/utils.go

@@ -2,9 +2,10 @@ package providers
 
 import (
 	"regexp"
-	"slices"
 	"strings"
 
+	"github.com/open-policy-agent/opa/v1/util"
+
 	"github.com/styrainc/regal/internal/lsp/cache"
 	"github.com/styrainc/regal/internal/lsp/types"
 )
@@ -47,15 +48,8 @@ func groupKeyedRefsByDepth(refs map[string]types.Ref) ([]int, map[int]map[string
 		byDepth[depth][key] = item
 	}
 
-	depths := make([]int, 0)
-	for k := range byDepth {
-		depths = append(depths, k)
-	}
-
 	// items from higher depths should be shown first
-	slices.Sort(depths)
-
-	return depths, byDepth
+	return util.KeysSorted(byDepth), byDepth
 }
 
 // inRuleBody is a best-effort helper to determine if the current line is in a rule body.

internal/lsp/completions/refs/used.go

@@ -14,11 +14,17 @@ import (
 	"github.com/styrainc/regal/pkg/builtins"
 	"github.com/styrainc/regal/pkg/config"
 
+	"github.com/styrainc/roast/pkg/rast"
 	"github.com/styrainc/roast/pkg/transform"
 
 	_ "embed"
 )
 
+var (
+	refNamesQuery = rast.RefStringToBody(`data.regal.lsp.completion.ref_names`)
+	pqOnce        = sync.OnceValues(prepareQuery)
+)
+
 // initialize prepares the rego query for finding ref names used in a module.
 // This is run and the resulting prepared query stored for performance reasons.
 // This function is only used by language server code paths and so init() is not
@@ -41,7 +47,7 @@ func prepareQuery() (*rego.PreparedEvalQuery, error) {
 	regoArgs := append([]func(*rego.Rego){
 		rego.ParsedBundle("regal", &rbundle.LoadedBundle),
 		rego.ParsedBundle("internal", &dataBundle),
-		rego.Query(`data.regal.lsp.completion.ref_names`),
+		rego.ParsedQuery(refNamesQuery),
 	}, builtins.RegalBuiltinRegoFuncs...)
 
 	preparedQuery, err := rego.New(regoArgs...).PrepareForEval(context.Background())
@@ -52,14 +58,12 @@ func prepareQuery() (*rego.PreparedEvalQuery, error) {
 	return &preparedQuery, nil
 }
 
-var pqOnce = sync.OnceValues(prepareQuery)
-
 // UsedInModule returns a list of ref names suitable for completion that are
 // used in the module's code.
 // See the rego above for more details on what's included and excluded.
 // This function is run when the parse completes for a module.
 func UsedInModule(ctx context.Context, module *ast.Module) ([]string, error) {
-	inputValue, err := transform.ToOPAInputValue(module)
+	inputValue, err := transform.ModuleToValue(module)
 	if err != nil {
 		return nil, fmt.Errorf("failed converting input to value: %w", err)
 	}

internal/lsp/handle_text_document_code_action_test.go

@@ -37,6 +37,10 @@ func TestHandleTextDocumentCodeAction(t *testing.T) {
 	params := types.CodeActionParams{
 		TextDocument: types.TextDocumentIdentifier{URI: uri},
 		Context:      types.CodeActionContext{Diagnostics: []types.Diagnostic{diag}},
+		Range: types.Range{
+			Start: types.Position{Line: 2, Character: 4},
+			End:   types.Position{Line: 2, Character: 10},
+		},
 	}
 
 	expectedAction := types.CodeAction{
@@ -115,8 +119,10 @@ func TestHandleTextDocumentCodeAction(t *testing.T) {
 	}
 }
 
-// 0.06 milliseconds per operation, not bad at all!
-// 63243 ns/op	   59576 B/op	    1110 allocs/op
+// 63243 ns/op	   59576 B/op	    1110 allocs/op - the OPA JSON roundtrip method
+// 42402 ns/op	   37822 B/op	     738 allocs/op - build input Value by hand
+// 45049 ns/op	   39731 B/op	     790 allocs/op - build input Value using reflection
+// 44024 ns/op	   38040 B/op	     749 allocs/op - build input Value using reflection + interning
 // ...
 // "real world" usage shows a number somewhere between 0.1 - 0.5 ms
 // of which most of the cost is in JSON marshaling and unmarshaling.

internal/lsp/hover/hover.go

@@ -165,14 +165,9 @@ func UpdateBuiltinPositions(cache *cache.Cache, uri string, builtins map[string]
 }
 
 func UpdateKeywordLocations(ctx context.Context, cache *cache.Cache, uri string) error {
-	module, ok := cache.GetModule(uri)
-	if !ok {
-		return fmt.Errorf("failed to update builtin positions: no parsed module for uri %q", uri)
-	}
-
-	fileContents, ok := cache.GetFileContents(uri)
+	fileContents, module, ok := cache.GetContentAndModule(uri)
 	if !ok {
-		return fmt.Errorf("failed to determine keyword locations: no file contents for uri %q", uri)
+		return fmt.Errorf("failed to determine keyword locations: missing file contents for uri %q", uri)
 	}
 
 	keywords, err := rego.AllKeywords(ctx, filepath.Base(uri), fileContents, module)