Add Debug Code Lens (#1103)

Following the excellent work on debugger support by @johanfylling,
this PR introduces a code lens for debugging, which triggers a
command that simply returns a launch configuration where the package
or rule clicked is set as the entrypoint.

The code lens provider is now also rewritten in Rego.

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

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)
+}

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}},
+		},
+	]
+}

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,
+	},
+}

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.