Use new setting to cache AST conversions from store

This doesn't have a huge impact on `regal lint` (although it does
result in less allocations), but could potentially make a real difference
for the LSP policies where ASTs from the whole workspace is provided under
`data.workspace`.

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

internal/lsp/completions/providers/policy.go

@@ -11,7 +11,6 @@ import (
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/rego"
 	"github.com/open-policy-agent/opa/storage"
-	"github.com/open-policy-agent/opa/topdown"
 
 	rbundle "github.com/styrainc/regal/bundle"
 	rio "github.com/styrainc/regal/internal/io"
@@ -139,13 +138,14 @@ func prepareQuery(ctx context.Context, store storage.Store, query string) (*rego
 
 func prepareRegoArgs(store storage.Store, query ast.Body) []func(*rego.Rego) {
 	return []func(*rego.Rego){
+		rego.StoreReadAST(true),
 		rego.Store(store),
 		rego.ParsedQuery(query),
 		rego.ParsedBundle("regal", &rbundle.LoadedBundle),
 		rego.Function2(builtins.RegalParseModuleMeta, builtins.RegalParseModule),
 		rego.Function1(builtins.RegalLastMeta, builtins.RegalLast),
-		// TODO: remove later
-		rego.EnablePrintStatements(true),
-		rego.PrintHook(topdown.NewPrintHook(os.Stderr)),
+		// Uncomment for development
+		// rego.EnablePrintStatements(true),
+		// rego.PrintHook(topdown.NewPrintHook(os.Stderr)),
 	}
 }

internal/lsp/rego/rego.go

@@ -118,6 +118,7 @@ func initialize() {
 			rego.ParsedBundle("regal", &rbundle.LoadedBundle),
 			rego.Function2(builtins.RegalParseModuleMeta, builtins.RegalParseModule),
 			rego.Function1(builtins.RegalLastMeta, builtins.RegalLast),
+			rego.StoreReadAST(true),
 		}, args...)
 	}
 

pkg/linter/linter.go

@@ -709,6 +709,7 @@ func (l Linter) prepareRegoArgs(query ast.Body) ([]func(*rego.Rego), error) {
 	}
 
 	regoArgs = append(regoArgs,
+		rego.StoreReadAST(true),
 		rego.Metrics(l.metrics),
 		rego.ParsedQuery(query),
 		rego.ParsedBundle("regal_eval_params", &dataBundle),

pkg/linter/linter_test.go

@@ -760,3 +760,17 @@ import data.unresolved`,
 		t.Fatalf("unexpected files: %v", got)
 	}
 }
+
+func BenchmarkRegalLintingItself(b *testing.B) {
+	linter := NewLinter().WithInputPaths([]string{"../../bundle"}).WithEnableAll(true)
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		_, err := linter.Lint(context.Background())
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}