Add regal eval:use-as-input directive (StyraInc#1088)

Placing `# regal eval:use-as-input` as the top comment in any
policy will now have the code lens eval feature automatically
use the roAST of that file as input, allowing you to query the
file as you work on it.

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

bundle/regal/ast/comments.rego

@@ -56,21 +56,37 @@ ignore_directives[row] := rules if {
 # METADATA
 # description: |
 #   returns an array of partitions, i.e. arrays containing all comments
-#   grouped by their "blocks".
-comment_blocks(comments) := [partition |
-	rows := [row |
-		some comment in comments
-		row := util.to_location_object(comment.location).row
+#   grouped by their "blocks". only comments on the same column as the
+#   one before is considered to be part of a block.
+comment_blocks(comments) := blocks if {
+	row_partitions := [partition |
+		rows := [row |
+			some comment in comments
+			row := util.to_location_object(comment.location).row
+		]
+		breaks := _splits(rows)
+
+		some j, k in breaks
+		partition := array.slice(
+			comments,
+			breaks[j - 1] + 1,
+			k + 1,
+		)
 	]
-	breaks := _splits(rows)
 
-	some j, k in breaks
-	partition := array.slice(
-		comments,
-		breaks[j - 1] + 1,
-		k + 1,
-	)
-]
+	blocks := [block |
+		some row_partition in row_partitions
+		some block in {col: partition |
+			some comment in row_partition
+			col := util.to_location_object(comment.location).col
+
+			partition := [c |
+				some c in row_partition
+				util.to_location_object(c.location).col == col
+			]
+		}
+	]
+}
 
 _splits(xs) := array.concat(
 	array.concat(

bundle/regal/rules/style/detached-metadata/detached_metadata_test.rego

@@ -53,7 +53,7 @@ allow := true
 }
 
 test_success_attached_metadata if {
-	r := rule.report with input as ast.policy(`
+	r := rule.report with input as ast.with_rego_v1(`
 # METADATA
 # title: valid
 allow := true
@@ -62,9 +62,7 @@ allow := true
 }
 
 test_success_detached_document_scope_ok if {
-	r := rule.report with input as regal.parse_module("p.rego", `
-package p
-
+	r := rule.report with input as ast.with_rego_v1(`
 # METADATA
 # scope: document
 # description: allow allows
@@ -75,3 +73,12 @@ allow := true
 `)
 	r == set()
 }
+
+test_success_not_detached_by_comment_in_different_column if {
+	r := rule.report with input as ast.with_rego_v1(`
+# METADATA
+# title: allow
+allow := true # not in block
+`)
+	r == set()
+}

docs/CONTRIBUTING.md

@@ -48,6 +48,23 @@ Regal.
 If you're struggling with any of the above points, or you're unsure of what to do, no worries! Just say so in your PR,
 or ask for advice in the `#regal` channel in the Styra Community [Slack](https://communityinviter.com/apps/styracommunity/signup)!
 
+### Rules Development Workflow
+
+Linter rules use a JSON representation of the AST as input. Use the `regal parse` command pointed at any Rego
+file to inspect its AST in this format. It is often more convenient to direct this output to a file, like
+`input.json` to browse it in your favorite editor, e.g. `regal parse policy.rego > input.json`.
+
+If you're using VS Code and the [OPA VS Code extension](https://github.com/open-policy-agent/vscode-opa), you may
+use the [Code Lens for Evaluation](https://docs.styra.com/regal/language-server#code-lenses-evaluation) to directly
+evaluate packages and rules using the `input.json` file as input, and see the result directly in your editor on the
+line you clicked to evaluate.
+
+As another convenience, any `.rego` file where the first comment in the policy is `# regal eval:use-as-input` will have
+the evaluation feature automatically use the AST of the file as input. This allows building queries against the AST of
+the policy you're working on, providing an extremely fast feedback loop for developing new rules!
+
+![Use AST of file as input](./assets/lsp/eval_use_as_input.png)
+
 ## Building
 
 Build the `regal` executable simply by running `go build`, or with `rq` installed, by running `build/do.rq build`.

docs/custom-rules.md

@@ -179,6 +179,23 @@ Starting from top to bottom, these are the components comprising our custom rule
    will later be included in the final report provided by Regal.
 1. The `result.location` helps extract the location from the element failing the test. Make sure to use it!
 
+### Rule Development Workflow
+
+In addition to making use of the `regal parse` command to inspect the AST of a policy, using Regal's
+[language server](https://docs.styra.com/regal/language-server) for rule development provides the absolute best rule
+development experience.
+
+If you're using VS Code and the [OPA VS Code extension](https://github.com/open-policy-agent/vscode-opa), you may
+use the [Code Lens for Evaluation](https://docs.styra.com/regal/language-server#code-lenses-evaluation) to directly
+evaluate packages and rules using the `input.json` file as input, and see the result directly in your editor on the
+line you clicked to evaluate.
+
+As another convenience, any `.rego` file where the first comment in the policy is `# regal eval:use-as-input` will have
+the evaluation feature automatically use the AST of the file as input. This allows building queries against the AST of
+the policy you're working on, providing an extremely fast feedback loop for developing new rules!
+
+![Use AST of file as input](./assets/lsp/eval_use_as_input.png)
+
 ## Aggregate Rules
 
 Aggregate rules are a special type of rule that allows you to collect data from multiple files before making a decision.

internal/lsp/server.go

@@ -2,6 +2,7 @@
 package lsp
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -409,6 +410,8 @@ func (l *LanguageServer) StartConfigWorker(ctx context.Context) {
 	}
 }
 
+var regalEvalUseAsInputComment = regexp.MustCompile(`^\s*regal eval:\s*use-as-input`)
+
 func (l *LanguageServer) StartCommandWorker(ctx context.Context) { // nolint:maintidx
 	// note, in this function conn.Call is used as the workspace/applyEdit message is a request, not a notification
 	// as per the spec. In order to be 'routed' to the correct handler on the client it must have an ID
@@ -545,7 +548,24 @@ func (l *LanguageServer) StartCommandWorker(ctx context.Context) { // nolint:mai
 				// if there are none, then it's a package evaluation
 				ruleHeadLocations := allRuleHeadLocations[path]
 
-				_, input := rio.FindInput(uri.ToPath(l.clientIdentifier, file), l.workspacePath())
+				var input io.Reader
+
+				// When the first comment in the file is `regal eval: use-as-input`, the AST of that module is
+				// used as the input rather than the contents of input.json. This is a development feature for
+				// working on rules (built-in or custom), allowing querying the AST of the module directly.
+				if regalEvalUseAsInputComment.Match(currentModule.Comments[0].Text) {
+					bs, err := encoding.JSON().Marshal(currentModule)
+					if err != nil {
+						l.logError(fmt.Errorf("failed to marshal module: %w", err))
+
+						break
+					}
+
+					input = bytes.NewReader(bs)
+				} else {
+					// Normal mode — try to find the input.json file in the workspace and use as input
+					_, input = rio.FindInput(uri.ToPath(l.clientIdentifier, file), l.workspacePath())
+				}
 
 				result, err := l.EvalWorkspacePath(ctx, path, input)
 				if err != nil {