From 5cb67eaa310ece0a4c670390d95862f9d363c474 Mon Sep 17 00:00:00 2001
From: Anders Eknert <anders@styra.com>
Date: Tue, 12 Nov 2024 13:37:39 +0100
Subject: [PATCH] Use new setting to cache AST conversions from store (#1256)

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 |  8 ++++----
 internal/lsp/rego/rego.go                    |  1 +
 pkg/linter/linter.go                         |  1 +
 pkg/linter/linter_test.go                    | 14 ++++++++++++++
 4 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/internal/lsp/completions/providers/policy.go b/internal/lsp/completions/providers/policy.go
index 6a849e20..6754ac10 100644
--- a/internal/lsp/completions/providers/policy.go
+++ b/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)),
 	}
 }
diff --git a/internal/lsp/rego/rego.go b/internal/lsp/rego/rego.go
index f113e29c..91163b21 100644
--- a/internal/lsp/rego/rego.go
+++ b/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...)
 	}
 
diff --git a/pkg/linter/linter.go b/pkg/linter/linter.go
index 1fb3d8db..9491c562 100644
--- a/pkg/linter/linter.go
+++ b/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),
diff --git a/pkg/linter/linter_test.go b/pkg/linter/linter_test.go
index ecce4bd6..7f54e1fe 100644
--- a/pkg/linter/linter_test.go
+++ b/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)
+		}
+	}
+}