diff --git a/bundle/regal/lsp/codelens/codelens.rego b/bundle/regal/lsp/codelens/codelens.rego new file mode 100644 index 00000000..8587b943 --- /dev/null +++ b/bundle/regal/lsp/codelens/codelens.rego @@ -0,0 +1,82 @@ +# METADATA +# description: | +# the code lens provider decides where code lenses should be placed in the given input file +# schemas: +# - input: schema.regal.ast +package regal.lsp.codelens + +import rego.v1 + +import data.regal.ast +import data.regal.result +import data.regal.util + +import data.regal.lsp.util.location + +# code lenses are displayed in the order they come back in the returned +# array, and 'evaluate' somehow feels better to the left of 'debug' +lenses := array.concat( + [l | some l in _eval_lenses], + [l | some l in _debug_lenses], +) + +_eval_lenses contains { + "range": location.to_range(result.ranged_location_from_text(input["package"]).location), + "command": { + "title": "Evaluate", + "command": "regal.eval", + "arguments": [ + input.regal.file.name, + ast.ref_to_string(input["package"].path), + util.to_location_object(input["package"].location).row, + ], + }, +} + +_eval_lenses contains _rule_lens(rule, "regal.eval", "Evaluate") if { + some rule in ast.rules +} + +_debug_lenses contains { + "range": location.to_range(result.ranged_location_from_text(input["package"]).location), + "command": { + "title": "Debug", + "command": "regal.debug", + "arguments": [ + input.regal.file.name, + ast.ref_to_string(input["package"].path), + util.to_location_object(input["package"].location).row, + ], + }, +} + +_debug_lenses contains _rule_lens(rule, "regal.debug", "Debug") if { + some rule in ast.rules + + # no need to add a debug lens for a rule like `pi := 3.14` + not _unconditional_constant(rule) +} + +_rule_lens(rule, command, title) := { + "range": location.to_range(result.ranged_location_from_text(rule).location), + "command": { + "title": title, + "command": command, + "arguments": [ + input.regal.file.name, # regal ignore:external-reference + sprintf("%s.%s", [ast.ref_to_string(input["package"].path), ast.ref_static_to_string(rule.head.ref)]), + util.to_location_object(rule.head.location).row, + ], + }, +} + +_rule_lens_args(filename, rule) := [ + filename, + sprintf("%s.%s", [ast.ref_to_string(input["package"].path), ast.ref_static_to_string(rule.head.ref)]), + util.to_location_object(rule.head.location).row, +] + +_unconditional_constant(rule) if { + not rule.body + ast.is_constant(rule.head.value) +} diff --git a/bundle/regal/lsp/codelens/codelens_test.rego b/bundle/regal/lsp/codelens/codelens_test.rego new file mode 100644 index 00000000..12f57623 --- /dev/null +++ b/bundle/regal/lsp/codelens/codelens_test.rego @@ -0,0 +1,61 @@ +package regal.lsp.codelens_test + +import rego.v1 + +import data.regal.lsp.codelens + +# regal ignore:rule-length +test_code_lenses_for_module if { + module := regal.parse_module("policy.rego", ` + package foo + + import rego.v1 + + rule1 := 1 + + rule2 if 1 + rule1 == 2 + `) + lenses := codelens.lenses with input as module + + lenses == [ + { + "command": { + "arguments": ["policy.rego", "data.foo", 2], + "command": "regal.eval", + "title": "Evaluate", + }, + "range": {"end": {"character": 8, "line": 1}, "start": {"character": 1, "line": 1}}, + }, + { + "command": { + "arguments": ["policy.rego", "data.foo.rule1", 6], + "command": "regal.eval", + "title": "Evaluate", + }, + "range": {"end": {"character": 11, "line": 5}, "start": {"character": 1, "line": 5}}, + }, + { + "command": { + "arguments": ["policy.rego", "data.foo.rule2", 8], + "command": "regal.eval", "title": "Evaluate", + }, + "range": {"end": {"character": 24, "line": 7}, "start": {"character": 1, "line": 7}}, + }, + { + "command": { + "arguments": ["policy.rego", "data.foo", 2], + "command": "regal.debug", + "title": "Debug", + }, + "range": {"end": {"character": 8, "line": 1}, "start": {"character": 1, "line": 1}}, + }, + { + "command": { + "arguments": ["policy.rego", "data.foo.rule2", 8], + "command": "regal.debug", + "title": "Debug", + }, + "range": {"end": {"character": 24, "line": 7}, "start": {"character": 1, "line": 7}}, + }, + ] +} diff --git a/bundle/regal/lsp/util/location/location.rego b/bundle/regal/lsp/util/location/location.rego new file mode 100644 index 00000000..c3c0a2ad --- /dev/null +++ b/bundle/regal/lsp/util/location/location.rego @@ -0,0 +1,16 @@ +package regal.lsp.util.location + +import rego.v1 + +# METADATA +# description: turns an AST location _with end attribute_ into an LSP range +to_range(location) := { + "start": { + "line": location.row - 1, + "character": location.col - 1, + }, + "end": { + "line": location.end.row - 1, + "character": location.end.col - 1, + }, +} diff --git a/internal/lsp/rego/rego.go b/internal/lsp/rego/rego.go index 4a53fc67..082f6eb9 100644 --- a/internal/lsp/rego/rego.go +++ b/internal/lsp/rego/rego.go @@ -98,14 +98,21 @@ func AllBuiltinCalls(module *ast.Module) []BuiltInCall { } //nolint:gochecknoglobals -var keywordsPreparedQuery *rego.PreparedEvalQuery - -//nolint:gochecknoglobals -var ruleHeadLocationsPreparedQuery *rego.PreparedEvalQuery +var ( + keywordsPreparedQuery *rego.PreparedEvalQuery + ruleHeadLocationsPreparedQuery *rego.PreparedEvalQuery + codeLensPreparedQuery *rego.PreparedEvalQuery +) //nolint:gochecknoglobals var preparedQueriesInitOnce sync.Once +type policy struct { + fileName string + contents string + module *ast.Module +} + func initialize() { regalRules := rio.MustLoadRegalBundleFS(rbundle.Bundle) @@ -134,74 +141,96 @@ func initialize() { } ruleHeadLocationsPreparedQuery = &rhlpq + + codeLensRegoArgs := createArgs(rego.Query("data.regal.lsp.codelens.lenses")) + + clpq, err := rego.New(codeLensRegoArgs...).PrepareForEval(context.Background()) + if err != nil { + panic(err) + } + + codeLensPreparedQuery = &clpq } // AllKeywords returns all keywords in the module. func AllKeywords(ctx context.Context, fileName, contents string, module *ast.Module) (map[string][]KeywordUse, error) { preparedQueriesInitOnce.Do(initialize) - enhancedInput, err := parse.PrepareAST(fileName, contents, module) + var keywords map[string][]KeywordUse + + value, err := queryToValue(ctx, keywordsPreparedQuery, policy{fileName, contents, module}, keywords) if err != nil { - return nil, fmt.Errorf("failed enhancing input: %w", err) + return nil, fmt.Errorf("failed querying code lenses: %w", err) } - rs, err := keywordsPreparedQuery.Eval(ctx, rego.EvalInput(enhancedInput)) + return value, nil +} + +// AllRuleHeadLocations returns mapping of rules names to the head locations. +func AllRuleHeadLocations(ctx context.Context, fileName, contents string, module *ast.Module) (RuleHeads, error) { + preparedQueriesInitOnce.Do(initialize) + + var ruleHeads RuleHeads + + value, err := queryToValue(ctx, ruleHeadLocationsPreparedQuery, policy{fileName, contents, module}, ruleHeads) if err != nil { - return nil, fmt.Errorf("failed evaluating keywords: %w", err) + return nil, fmt.Errorf("failed querying code lenses: %w", err) } - if len(rs) != 1 { - return nil, errors.New("expected exactly one result from evaluation") - } + return value, nil +} - if len(rs[0].Expressions) != 1 { - return nil, errors.New("expected exactly one expression in result") - } +// CodeLenses returns all code lenses in the module. +func CodeLenses(ctx context.Context, uri, contents string, module *ast.Module) ([]types.CodeLens, error) { + preparedQueriesInitOnce.Do(initialize) - var result map[string][]KeywordUse + var codeLenses []types.CodeLens - err = rio.JSONRoundTrip(rs[0].Expressions[0].Value, &result) + value, err := queryToValue(ctx, codeLensPreparedQuery, policy{uri, contents, module}, codeLenses) if err != nil { - return nil, fmt.Errorf("failed unmarshaling keywords: %w", err) + return nil, fmt.Errorf("failed querying code lenses: %w", err) } - return result, nil + return value, nil } -// AllRuleHeadLocations returns mapping of rules names to the head locations. -func AllRuleHeadLocations(ctx context.Context, fileName, contents string, module *ast.Module) (RuleHeads, error) { - preparedQueriesInitOnce.Do(initialize) +func queryToValue[T any](ctx context.Context, pq *rego.PreparedEvalQuery, policy policy, toValue T) (T, error) { + input, err := parse.PrepareAST(policy.fileName, policy.contents, policy.module) + if err != nil { + return toValue, fmt.Errorf("failed to prepare input: %w", err) + } - enhancedInput, err := parse.PrepareAST(fileName, contents, module) + result, err := toValidResult(pq.Eval(ctx, rego.EvalInput(input))) if err != nil { - return nil, fmt.Errorf("failed enhancing input: %w", err) + return toValue, err //nolint:wrapcheck } - rs, err := ruleHeadLocationsPreparedQuery.Eval(ctx, rego.EvalInput(enhancedInput)) + err = rio.JSONRoundTrip(result.Expressions[0].Value, &toValue) if err != nil { - return nil, fmt.Errorf("failed evaluating keywords: %w", err) + return toValue, fmt.Errorf("failed unmarshaling code lenses: %w", err) + } + + return toValue, nil +} + +func toValidResult(rs rego.ResultSet, err error) (rego.Result, error) { + if err != nil { + return rego.Result{}, fmt.Errorf("evaluation failed: %w", err) } if len(rs) == 0 { - return nil, errors.New("no results returned from evaluation") + return rego.Result{}, errors.New("no results returned from evaluation") } if len(rs) != 1 { - return nil, errors.New("expected exactly one result from evaluation") + return rego.Result{}, errors.New("expected exactly one result from evaluation") } if len(rs[0].Expressions) != 1 { - return nil, errors.New("expected exactly one expression in result") - } - - var result RuleHeads - - err = rio.JSONRoundTrip(rs[0].Expressions[0].Value, &result) - if err != nil { - return nil, fmt.Errorf("failed unmarshaling keywords: %w", err) + return rego.Result{}, errors.New("expected exactly one expression in result") } - return result, nil + return rs[0], nil } // ToInput prepares a module with Regal additions to be used as input for evaluation. diff --git a/internal/lsp/rego/rego_test.go b/internal/lsp/rego/rego_test.go new file mode 100644 index 00000000..2a1ff70f --- /dev/null +++ b/internal/lsp/rego/rego_test.go @@ -0,0 +1,101 @@ +package rego + +import ( + "context" + "testing" + + "github.com/styrainc/regal/internal/parse" +) + +func TestCodeLenses(t *testing.T) { + t.Parallel() + + contents := `package p + + import rego.v1 + + allow if "foo" in input.bar` + + module := parse.MustParseModule(contents) + + lenses, er := CodeLenses(context.TODO(), "p.rego", contents, module) + if er != nil { + t.Fatalf("unexpected error: %v", er) + } + + // 2 for the package, 2 for the rule + // the contents of the lenses are tested in Rego + if len(lenses) != 4 { + t.Fatalf("expected 4 code lenses, got %d", len(lenses)) + } +} + +func TestAllRuleHeadLocations(t *testing.T) { + t.Parallel() + + contents := `package p + + import rego.v1 + + default allow := false + + allow if 1 + allow if 2 + + foo.bar[x] if x := 1 + foo.bar[x] if x := 2` + + module := parse.MustParseModule(contents) + + ruleHeads, er := AllRuleHeadLocations(context.TODO(), "p.rego", contents, module) + if er != nil { + t.Fatalf("unexpected error: %v", er) + } + + if len(ruleHeads) != 2 { + t.Fatalf("expected 2 code lenses, got %d", len(ruleHeads)) + } + + if len(ruleHeads["data.p.allow"]) != 3 { + t.Fatalf("expected 3 allow rule heads, got %d", len(ruleHeads["data.p.allow"])) + } + + if len(ruleHeads["data.p.foo.bar"]) != 2 { + t.Fatalf("expected 2 foo.bar rule heads, got %d", len(ruleHeads["data.p.foo.bar"])) + } +} + +func TestAllKeywords(t *testing.T) { + t.Parallel() + + contents := `package p + + import rego.v1 + + my_set contains "x" if true + ` + + module := parse.MustParseModule(contents) + + keywords, er := AllKeywords(context.TODO(), "p.rego", contents, module) + if er != nil { + t.Fatalf("unexpected error: %v", er) + } + + // this is "lines with keywords", not number of keywords + if len(keywords) != 3 { + t.Fatalf("expected 1 keyword, got %d", len(keywords)) + } + + if len(keywords["1"]) != 1 { + t.Fatalf("expected 1 keywords on line 1, got %d", len(keywords["1"])) + } + + if len(keywords["3"]) != 1 { + t.Fatalf("expected 1 keywords on line 3, got %d", len(keywords["1"])) + } + + if len(keywords["5"]) != 2 { + t.Fatalf("expected 2 keywords on line 5, got %d", len(keywords["1"])) + } +} diff --git a/internal/lsp/server.go b/internal/lsp/server.go index ccc98dd7..fe52e119 100644 --- a/internal/lsp/server.go +++ b/internal/lsp/server.go @@ -544,6 +544,52 @@ func (l *LanguageServer) StartCommandWorker(ctx context.Context) { // nolint:mai // handle this ourselves as it's a rename and not a content edit fixed = false + case "regal.debug": + if l.clientIdentifier != clients.IdentifierVSCode { + l.logError(errors.New("regal.debug command is only supported in VSCode")) + + break + } + + if len(params.Arguments) != 3 { + l.logError(fmt.Errorf("expected three arguments, got %d", len(params.Arguments))) + + break + } + + file, ok := params.Arguments[0].(string) + if !ok { + l.logError(fmt.Errorf("expected first argument to be a string, got %T", params.Arguments[0])) + + break + } + + path, ok := params.Arguments[1].(string) + if !ok { + l.logError(fmt.Errorf("expected second argument to be a string, got %T", params.Arguments[1])) + + break + } + + inputPath, _ := rio.FindInput(uri.ToPath(l.clientIdentifier, file), l.workspacePath()) + + responseParams := map[string]any{ + "type": "opa-debug", + "name": "Debug " + path, + "request": "launch", + "command": "eval", + "query": path, + "enablePrint": true, + "stopOnEntry": true, + "inputPath": inputPath, + } + + responseResult := map[string]any{} + + err = l.conn.Call(ctx, "regal/startDebugging", responseParams, &responseResult) + if err != nil { + l.logError(fmt.Errorf("regal/startDebugging failed: %v", err.Error())) + } case "regal.eval": if len(params.Arguments) != 3 { l.logError(fmt.Errorf("expected three arguments, got %d", len(params.Arguments))) @@ -1347,7 +1393,6 @@ func (l *LanguageServer) handleWorkspaceExecuteCommand( req *jsonrpc2.Request, ) (result any, err error) { var params types.ExecuteCommandParams - if err := encoding.JSON().Unmarshal(*req.Params, ¶ms); err != nil { return nil, fmt.Errorf("failed to unmarshal params: %w", err) } @@ -1403,7 +1448,7 @@ func (l *LanguageServer) handleTextDocumentInlayHint( } func (l *LanguageServer) handleTextDocumentCodeLens( - _ context.Context, + ctx context.Context, _ *jsonrpc2.Conn, req *jsonrpc2.Request, ) (result any, err error) { @@ -1414,68 +1459,15 @@ func (l *LanguageServer) handleTextDocumentCodeLens( module, ok := l.cache.GetModule(params.TextDocument.URI) if !ok { - // return a null response, as per the spec - return nil, nil - } - - codeLenses := make([]types.CodeLens, 0) - - // Package - - pkgLens := types.CodeLens{ - Range: locationToRange(module.Package.Location), - Command: &types.Command{ - Title: "Evaluate", - Command: "regal.eval", - Arguments: &[]any{ - module.Package.Location.File, - module.Package.Path.String(), - module.Package.Location.Row, - }, - }, + return nil, nil // return a null response, as per the spec } - codeLenses = append(codeLenses, pkgLens) - - // Rules - - for _, rule := range module.Rules { - if rule.Head.Args != nil { - // Skip functions for now, as it's not clear how to best - // provide inputs for them. - continue - } - - ruleLens := types.CodeLens{ - Range: locationToRange(rule.Location), - Command: &types.Command{ - Title: "Evaluate", - Command: "regal.eval", - Arguments: &[]any{ - module.Package.Location.File, - module.Package.Path.String() + "." + getRuleName(rule), - rule.Head.Location.Row, - }, - }, - } - - codeLenses = append(codeLenses, ruleLens) - } - - return codeLenses, nil -} - -func getRuleName(rule *ast.Rule) string { - result := rule.Head.Ref().String() - - // only evaluate the top level rule name if there are refs - // e.g. my[foo].bar -> my - // bar.bar.bar -> bar.bar.bar - if i := strings.Index(result, "["); i > 0 { - return result[:i] + contents, ok := l.cache.GetFileContents(params.TextDocument.URI) + if !ok { + return nil, nil // return a null response, as per the spec } - return result + return rego.CodeLenses(ctx, params.TextDocument.URI, contents, module) //nolint:wrapcheck } func (l *LanguageServer) handleTextDocumentCompletion( @@ -2196,6 +2188,7 @@ func (l *LanguageServer) handleInitialize( }, ExecuteCommandProvider: types.ExecuteCommandOptions{ Commands: []string{ + "regal.debug", "regal.eval", "regal.fix.opa-fmt", "regal.fix.use-rego-v1", diff --git a/pkg/builtins/builtins.go b/pkg/builtins/builtins.go index 5ad7dcd0..77e32a59 100644 --- a/pkg/builtins/builtins.go +++ b/pkg/builtins/builtins.go @@ -2,7 +2,6 @@ package builtins import ( - "bytes" "errors" "github.com/anderseknert/roast/pkg/encoding" @@ -62,14 +61,12 @@ func RegalParseModule(_ rego.BuiltinContext, filename *ast.Term, policy *ast.Ter return nil, err } - json := encoding.JSON() - - var buf bytes.Buffer - if err := json.NewEncoder(&buf).Encode(&enhancedAST); err != nil { + roast, err := encoding.JSON().MarshalToString(enhancedAST) + if err != nil { return nil, err } - term, err := ast.ParseTerm(buf.String()) + term, err := ast.ParseTerm(roast) if err != nil { return nil, err }