lsp: Add code lens support for evaluating rules (#968)

- This is currently limited to VS Code clients, but may be extended
  to other clients too in the furure, if they can support it on their
  side.

- Implementation wil recursively look for an `input.json` file and
  pick the one closest to the file evaluated, down to the level of
  the workspace root (if any file is found).

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

README.md

@@ -702,7 +702,7 @@ in the near future:
 
 - [ ] Make "Check on save" unnecessary by allowing diagnostics to include
       [compilation errors](https://github.com/StyraInc/regal/issues/745)
-- [ ] Add Code Lens to "Evaluate" any rule or package (VS Code only, initially)
+- [x] Add Code Lens to "Evaluate" any rule or package (VS Code only, initially)
 - [ ] Implement [Signature Help](https://github.com/StyraInc/regal/issues/695) feature
 
 The roadmap is updated when all the current items have been completed.

bundle/regal/ast/search.rego

@@ -287,11 +287,6 @@ find_some_decl_names_in_scope(rule, location) := {some_var.value |
 	_before_location(rule, some_var, location)
 }
 
-# _rules_with_bodies[rule_index] := rule if {
-# 	some rule_index, rule in input.rules
-# 	not generated_body(rule)
-# }
-
 exprs[rule_index][expr_index] := expr if {
 	some rule_index, rule in input.rules
 	some expr_index, expr in rule.body

bundle/regal/main.rego

@@ -12,10 +12,10 @@ lint.aggregates := aggregate
 
 lint.ignore_directives[input.regal.file.name] := ast.ignore_directives
 
-lint.violations := report
-
 lint_aggregate.violations := aggregate_report
 
+lint.violations := report
+
 rules_to_run[category][title] if {
 	some category, title
 	config.merged_config.rules[category][title]

internal/lsp/eval.go

@@ -0,0 +1,113 @@
+package lsp
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"path"
+	"path/filepath"
+	"strings"
+
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/bundle"
+	"github.com/open-policy-agent/opa/rego"
+	"github.com/open-policy-agent/opa/topdown"
+
+	"github.com/styrainc/regal/internal/lsp/uri"
+	"github.com/styrainc/regal/pkg/builtins"
+)
+
+func (l *LanguageServer) Eval(ctx context.Context, query string, input string) (rego.ResultSet, error) {
+	modules := l.cache.GetAllModules()
+	moduleFiles := make([]bundle.ModuleFile, 0, len(modules))
+
+	for fileURI, module := range modules {
+		moduleFiles = append(moduleFiles, bundle.ModuleFile{
+			URL:    fileURI,
+			Parsed: module,
+			Path:   uri.ToPath(l.clientIdentifier, fileURI),
+		})
+	}
+
+	bd := bundle.Bundle{
+		Data: make(map[string]any),
+		Manifest: bundle.Manifest{
+			Roots:    &[]string{""},
+			Metadata: map[string]any{"name": "workspace"},
+		},
+		Modules: moduleFiles,
+	}
+
+	regoArgs := prepareRegoArgs(ast.MustParseBody(query), bd)
+
+	pq, err := rego.New(regoArgs...).PrepareForEval(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed preparing query: %w", err)
+	}
+
+	if input != "" {
+		inputMap := make(map[string]any)
+
+		err = json.Unmarshal([]byte(input), &inputMap)
+		if err != nil {
+			return nil, fmt.Errorf("failed unmarshalling input: %w", err)
+		}
+
+		return pq.Eval(ctx, rego.EvalInput(inputMap)) //nolint:wrapcheck
+	}
+
+	return pq.Eval(ctx) //nolint:wrapcheck
+}
+
+type EvalPathResult struct {
+	Value       any  `json:"value"`
+	IsUndefined bool `json:"isUndefined"`
+}
+
+func FindInput(file string, workspacePath string) string {
+	relative := strings.TrimPrefix(file, workspacePath)
+	components := strings.Split(path.Dir(relative), string(filepath.Separator))
+
+	for i := range len(components) {
+		inputPath := path.Join(workspacePath, path.Join(components[:len(components)-i]...), "input.json")
+
+		if input, err := os.ReadFile(inputPath); err == nil {
+			return string(input)
+		}
+	}
+
+	return ""
+}
+
+func (l *LanguageServer) EvalWorkspacePath(ctx context.Context, query string, input string) (EvalPathResult, error) {
+	resultQuery := "result := " + query
+
+	result, err := l.Eval(ctx, resultQuery, input)
+	if err != nil {
+		return EvalPathResult{}, fmt.Errorf("failed evaluating query: %w", err)
+	}
+
+	if len(result) == 0 {
+		return EvalPathResult{IsUndefined: true}, nil
+	}
+
+	res, ok := result[0].Bindings["result"]
+	if !ok {
+		return EvalPathResult{}, errors.New("expected result in bindings, didn't get it")
+	}
+
+	return EvalPathResult{Value: res}, nil
+}
+
+func prepareRegoArgs(query ast.Body, bd bundle.Bundle) []func(*rego.Rego) {
+	return []func(*rego.Rego){
+		rego.ParsedQuery(query),
+		rego.ParsedBundle("workspace", &bd),
+		rego.Function2(builtins.RegalParseModuleMeta, builtins.RegalParseModule),
+		rego.Function1(builtins.RegalLastMeta, builtins.RegalLast),
+		rego.EnablePrintStatements(true),
+		rego.PrintHook(topdown.NewPrintHook(os.Stderr)),
+	}
+}

internal/lsp/eval_test.go

@@ -0,0 +1,108 @@
+package lsp
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/styrainc/regal/internal/parse"
+)
+
+func TestEvalWorkspacePath(t *testing.T) {
+	t.Parallel()
+
+	ls := NewLanguageServer(&LanguageServerOptions{ErrorLog: os.Stderr})
+
+	policy1 := `package policy1
+
+	import rego.v1
+
+	import data.policy2
+
+	default allow := false
+
+	allow if policy2.allow
+	`
+
+	policy2 := `package policy2
+
+	import rego.v1
+
+	allow if input.exists
+	`
+
+	module1 := parse.MustParseModule(policy1)
+	module2 := parse.MustParseModule(policy2)
+
+	ls.cache.SetFileContents("file://policy1.rego", policy1)
+	ls.cache.SetFileContents("file://policy2.rego", policy2)
+	ls.cache.SetModule("file://policy1.rego", module1)
+	ls.cache.SetModule("file://policy2.rego", module2)
+
+	res, err := ls.EvalWorkspacePath(context.TODO(), "data.policy1.allow", `{"exists": true}`)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	empty := EvalPathResult{}
+
+	if res == empty {
+		t.Fatal("expected result, got nil")
+	}
+
+	if val, ok := res.Value.(bool); !ok || val != true {
+		t.Fatalf("expected true, got false")
+	}
+}
+
+func TestFindInput(t *testing.T) {
+	t.Parallel()
+
+	tmpDir := t.TempDir()
+
+	workspacePath := tmpDir + "/workspace"
+	file := tmpDir + "/workspace/foo/bar/baz.rego"
+
+	if err := os.MkdirAll(workspacePath+"/foo/bar", 0o755); err != nil {
+		t.Fatal(err)
+	}
+
+	if FindInput(file, workspacePath) != "" {
+		t.Fatalf("did not expect to find input.json")
+	}
+
+	content := `{"x": 1}`
+
+	createWithContent(t, tmpDir+"/workspace/foo/bar/input.json", content)
+
+	if res := FindInput(file, workspacePath); res != content {
+		t.Errorf("expected input at %s, got %s", content, res)
+	}
+
+	err := os.Remove(tmpDir + "/workspace/foo/bar/input.json")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	createWithContent(t, tmpDir+"/workspace/input.json", content)
+
+	if res := FindInput(file, workspacePath); res != content {
+		t.Errorf("expected input at %s, got %s", content, res)
+	}
+}
+
+func createWithContent(t *testing.T, path string, content string) {
+	t.Helper()
+
+	f, err := os.Create(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer f.Close()
+
+	_, err = f.WriteString(content)
+	if err != nil {
+		t.Fatal(err)
+	}
+}