Add source action to explore compiler stages

README.md

@@ -3,6 +3,7 @@
 [![Build Status](https://github.com/styrainc/regal/workflows/Build/badge.svg?branch=main)](https://github.com/styrainc/regal/actions)
 ![OPA v0.68.0](https://openpolicyagent.org/badge/v0.68.0)
 [![codecov](https://codecov.io/github/StyraInc/regal/graph/badge.svg?token=EQK01YF3X3)](https://codecov.io/github/StyraInc/regal)
+[![Downloads](https://img.shields.io/github/downloads/styrainc/regal/total.svg)](https://github.com/StyraInc/regal/releases)
 
 Regal is a linter and language server for [Rego](https://www.openpolicyagent.org/docs/latest/policy-language/), making
 your Rego magnificent, and you the ruler of rules!

cmd/languageserver.go

@@ -46,6 +46,7 @@ func init() {
 			go ls.StartConfigWorker(ctx)
 			go ls.StartWorkspaceStateWorker(ctx)
 			go ls.StartTemplateWorker(ctx)
+			go ls.StartWebServer(ctx)
 
 			sigChan := make(chan os.Signal, 1)
 			signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

docs/language-server.md

@@ -131,14 +131,20 @@ that appears on the line with a diagnostic message, or by pressing `ctrl/cmd + .
   src={require('./assets/lsp/codeaction.png').default}
   alt="Screenshot of code action displayed in Zed"/>
 
-Regal currently provides quick fix code actions for the following linter rules:
+Regal currently provides **quick fix actions** for the following linter rules:
 
 - [opa-fmt](https://docs.styra.com/regal/rules/style/opa-fmt)
 - [use-rego-v1](https://docs.styra.com/regal/rules/imports/use-rego-v1)
 - [use-assignment-operator](https://docs.styra.com/regal/rules/style/use-assignment-operator)
 - [no-whitespace-comment](https://docs.styra.com/regal/rules/style/no-whitespace-comment)
 - [directory-package-mismatch](https://docs.styra.com/regal/rules/idiomatic/directory-package-mismatch)
 
+Regal also provides **source actions** — actions that apply to a whole file and aren't triggered by linter issues:
+
+- **Explore compiler stages for policy** — Opens a browser window with an embedded version of the
+  [opa-explorer](https://github.com/srenatus/opa-explorer), where advanced users can explore the different stages
+  of the Rego compiler's output for a given policy.
+
 ### Code lenses (Evaluation)
 
 The code lens feature provides language servers a way to add actionable commands just next to the code that the action

internal/explorer/stages.go

@@ -0,0 +1,163 @@
+package explorer
+
+import (
+	"bytes"
+	"context"
+
+	"github.com/anderseknert/roast/pkg/encoding"
+
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/bundle"
+	"github.com/open-policy-agent/opa/compile"
+	"github.com/open-policy-agent/opa/ir"
+
+	compile2 "github.com/styrainc/regal/internal/compile"
+)
+
+type CompileResult struct {
+	Stage  string
+	Result *ast.Module
+	Error  string
+}
+
+type stage struct{ name, metricName string }
+
+// NOTE(sr): copied from 0.68.0.
+//
+//nolint:gochecknoglobals
+var stages []stage = []stage{
+	{"ResolveRefs", "compile_stage_resolve_refs"},
+	{"InitLocalVarGen", "compile_stage_init_local_var_gen"},
+	{"RewriteRuleHeadRefs", "compile_stage_rewrite_rule_head_refs"},
+	{"CheckKeywordOverrides", "compile_stage_check_keyword_overrides"},
+	{"CheckDuplicateImports", "compile_stage_check_duplicate_imports"},
+	{"RemoveImports", "compile_stage_remove_imports"},
+	{"SetModuleTree", "compile_stage_set_module_tree"},
+	{"SetRuleTree", "compile_stage_set_rule_tree"},
+	{"RewriteLocalVars", "compile_stage_rewrite_local_vars"},
+	{"CheckVoidCalls", "compile_stage_check_void_calls"},
+	{"RewritePrintCalls", "compile_stage_rewrite_print_calls"},
+	{"RewriteExprTerms", "compile_stage_rewrite_expr_terms"},
+	{"ParseMetadataBlocks", "compile_stage_parse_metadata_blocks"},
+	{"SetAnnotationSet", "compile_stage_set_annotationset"},
+	{"RewriteRegoMetadataCalls", "compile_stage_rewrite_rego_metadata_calls"},
+	{"SetGraph", "compile_stage_set_graph"},
+	{"RewriteComprehensionTerms", "compile_stage_rewrite_comprehension_terms"},
+	{"RewriteRefsInHead", "compile_stage_rewrite_refs_in_head"},
+	{"RewriteWithValues", "compile_stage_rewrite_with_values"},
+	{"CheckRuleConflicts", "compile_stage_check_rule_conflicts"},
+	{"CheckUndefinedFuncs", "compile_stage_check_undefined_funcs"},
+	{"CheckSafetyRuleHeads", "compile_stage_check_safety_rule_heads"},
+	{"CheckSafetyRuleBodies", "compile_stage_check_safety_rule_bodies"},
+	{"RewriteEquals", "compile_stage_rewrite_equals"},
+	{"RewriteDynamicTerms", "compile_stage_rewrite_dynamic_terms"},
+	{"RewriteTestRulesForTracing", "compile_stage_rewrite_test_rules_for_tracing"}, // must run after RewriteDynamicTerms
+	{"CheckRecursion", "compile_stage_check_recursion"},
+	{"CheckTypes", "compile_stage_check_types"},
+	{"CheckUnsafeBuiltins", "compile_state_check_unsafe_builtins"},
+	{"CheckDeprecatedBuiltins", "compile_state_check_deprecated_builtins"},
+	{"BuildRuleIndices", "compile_stage_rebuild_indices"},
+	{"BuildComprehensionIndices", "compile_stage_rebuild_comprehension_indices"},
+	{"BuildRequiredCapabilities", "compile_stage_build_required_capabilities"},
+}
+
+func CompilerStages(path, rego string, useStrict, useAnno, usePrint bool) []CompileResult {
+	c := compile2.NewCompilerWithRegalBuiltins().
+		WithStrict(useStrict).
+		WithEnablePrintStatements(usePrint).
+		WithUseTypeCheckAnnotations(useAnno)
+
+	result := make([]CompileResult, 0, len(stages)+1)
+	result = append(result, CompileResult{
+		Stage: "ParseModule",
+	})
+
+	mod, err := ast.ParseModuleWithOpts(path, rego, ast.ParserOptions{ProcessAnnotation: useAnno})
+	if err != nil {
+		result[0].Error = err.Error()
+
+		return result
+	}
+
+	result[0].Result = mod
+
+	for i := range stages {
+		stage := stages[i]
+		c = c.WithStageAfter(stage.name,
+			ast.CompilerStageDefinition{
+				Name:       stage.name + "Record",
+				MetricName: stage.metricName + "_record",
+				Stage: func(c0 *ast.Compiler) *ast.Error {
+					result = append(result, CompileResult{
+						Stage:  stage.name,
+						Result: getOne(c0.Modules),
+					})
+
+					return nil
+				},
+			})
+	}
+
+	c.Compile(map[string]*ast.Module{
+		path: mod,
+	})
+
+	if len(c.Errors) > 0 {
+		// stage after the last than ran successfully
+		stage := stages[len(result)-1]
+		result = append(result, CompileResult{
+			Stage: stage.name + ": Failure",
+			Error: c.Errors.Error(),
+		})
+	}
+
+	return result
+}
+
+func getOne(mods map[string]*ast.Module) *ast.Module {
+	for _, m := range mods {
+		return m.Copy()
+	}
+
+	panic("unreachable")
+}
+
+func Plan(ctx context.Context, path, rego string, usePrint bool) (string, error) {
+	mod, err := ast.ParseModuleWithOpts(path, rego, ast.ParserOptions{ProcessAnnotation: true})
+	if err != nil {
+		return "", err //nolint:wrapcheck
+	}
+
+	b := &bundle.Bundle{
+		Modules: []bundle.ModuleFile{
+			{
+				URL:    "/url",
+				Path:   path,
+				Raw:    []byte(rego),
+				Parsed: mod,
+			},
+		},
+	}
+
+	compiler := compile.New().
+		WithTarget(compile.TargetPlan).
+		WithBundle(b).
+		WithRegoAnnotationEntrypoints(true).
+		WithEnablePrintStatements(usePrint)
+	if err := compiler.Build(ctx); err != nil {
+		return "", err //nolint:wrapcheck
+	}
+
+	var policy ir.Policy
+
+	if err := encoding.JSON().Unmarshal(compiler.Bundle().PlanModules[0].Raw, &policy); err != nil {
+		return "", err //nolint:wrapcheck
+	}
+
+	buf := bytes.Buffer{}
+	if err := ir.Pretty(&buf, &policy); err != nil {
+		return "", err //nolint:wrapcheck
+	}
+
+	return buf.String(), nil
+}

internal/lsp/server.go

@@ -43,6 +43,7 @@ import (
 	rparse "github.com/styrainc/regal/internal/parse"
 	"github.com/styrainc/regal/internal/update"
 	"github.com/styrainc/regal/internal/util"
+	"github.com/styrainc/regal/internal/web"
 	"github.com/styrainc/regal/pkg/config"
 	"github.com/styrainc/regal/pkg/fixer"
 	"github.com/styrainc/regal/pkg/fixer/fileprovider"
@@ -81,6 +82,7 @@ func NewLanguageServer(opts *LanguageServerOptions) *LanguageServer {
 		templateFile:               make(chan fileUpdateEvent, 10),
 		configWatcher:              lsconfig.NewWatcher(&lsconfig.WatcherOpts{ErrorWriter: opts.ErrorLog}),
 		completionsManager:         completions.NewDefaultManager(c, store),
+		webServer:                  web.NewServer(c),
 	}
 
 	return ls
@@ -111,6 +113,8 @@ type LanguageServer struct {
 	builtinsPositionFile       chan fileUpdateEvent
 	commandRequest             chan types.ExecuteCommandParams
 	templateFile               chan fileUpdateEvent
+
+	webServer *web.Server
 }
 
 // fileUpdateEvent is sent to a channel when an update is required for a file.
@@ -823,6 +827,10 @@ func (l *LanguageServer) StartTemplateWorker(ctx context.Context) {
 	}
 }
 
+func (l *LanguageServer) StartWebServer(ctx context.Context) {
+	l.webServer.Start(ctx)
+}
+
 func (l *LanguageServer) templateContentsForFile(fileURI string) (string, error) {
 	content, ok := l.cache.GetFileContents(fileURI)
 	if !ok {
@@ -1225,6 +1233,23 @@ func (l *LanguageServer) handleTextDocumentCodeAction(
 		return actions, nil
 	}
 
+	// only VS Code has the capability to open a provided URL, as far as we know
+	// if we learn about others with this capability later, we should add them!
+	if l.clientIdentifier == clients.IdentifierVSCode {
+		explorerURL := l.webServer.GetBaseURL() + "/explorer" +
+			strings.TrimPrefix(params.TextDocument.URI, l.workspaceRootURI)
+
+		actions = append(actions, types.CodeAction{
+			Title: "Explore compiler stages for this policy",
+			Kind:  "source.explore",
+			Command: types.Command{
+				Title:     "Explore compiler stages for this policy",
+				Command:   "vscode.open",
+				Arguments: &[]any{explorerURL},
+			},
+		})
+	}
+
 	for _, diag := range params.Context.Diagnostics {
 		switch diag.Code {
 		case ruleNameOPAFmt:
@@ -2094,6 +2119,8 @@ func (l *LanguageServer) handleInitialize(
 		)
 	}
 
+	l.webServer.SetClient(l.clientIdentifier)
+
 	if params.InitializationOptions != nil {
 		l.clientInitializationOptions = *params.InitializationOptions
 	}
@@ -2137,7 +2164,10 @@ func (l *LanguageServer) handleInitialize(
 			},
 			HoverProvider: true,
 			CodeActionProvider: types.CodeActionOptions{
-				CodeActionKinds: []string{"quickfix"},
+				CodeActionKinds: []string{
+					"quickfix",
+					"source.explore",
+				},
 			},
 			ExecuteCommandProvider: types.ExecuteCommandOptions{
 				Commands: []string{
@@ -2182,6 +2212,8 @@ func (l *LanguageServer) handleInitialize(
 			return nil, fmt.Errorf("failed to load workspace contents: %w", err)
 		}
 
+		l.webServer.SetWorkspaceURI(l.workspaceRootURI)
+
 		l.diagnosticRequestWorkspace <- "server initialize"
 	}
 

internal/lsp/types/types.go

@@ -185,7 +185,7 @@ type CodeActionContext struct {
 type CodeAction struct {
 	Title       string       `json:"title"`
 	Kind        string       `json:"kind"`
-	Diagnostics []Diagnostic `json:"diagnostics"`
+	Diagnostics []Diagnostic `json:"diagnostics,omitempty"`
 	IsPreferred *bool        `json:"isPreferred,omitempty"`
 	Command     Command      `json:"command"`
 }

internal/web/assets/main.tpl

@@ -0,0 +1,33 @@
+<!DOCTYPE html>
+<html>
+<head>
+  <meta charset="utf-8"/>
+  <link rel="stylesheet" href="/assets/missing.min.css">
+  <script src="/assets/htmx-1.8.4.min.js"></script>
+</head>
+<body>
+<main class="crowded">
+  <form>
+    <div class="f-row">
+					<textarea name="code"
+                    style="height: 300px"
+                    hx-post="/?tmpl=output"
+                    hx-target="#output"
+                    hx-trigger="keyup changed delay:200ms"
+                    class="flex-grow:1 monospace"
+          >{{ .Code }}</textarea>
+    </div>
+  </form>
+  <section id="output">
+    {{ block "output" . }}
+    {{ range .Result }}
+    <details class={{ .Class }} {{ if .Show }}open{{ end }}>
+      <summary>{{ .Stage }}</summary>
+      <pre><code>{{ .Output }}</code></pre>
+    </details>
+    {{ end }}
+    {{ end }}
+  </section>
+</main>
+</body>
+</html>