From dff78c48b90134f4369c9571cfa665fc72f1b222 Mon Sep 17 00:00:00 2001
From: Patrick East <east.patrick@gmail.com>
Date: Tue, 19 May 2020 17:48:23 -0700
Subject: [PATCH] cover: Disable local var plugging on trace

Similar treatment as what we do with the `Profiler`, the coverage
tracer will now signal via configuration that it doesn't need local
variable metadata.

This improves the speed of evaluations with only coverage enabled by
a pretty substantial amount, in particular with many/larger variables.

Signed-off-by: Patrick East <east.patrick@gmail.com>
---
 cover/cover.go            |  7 ++++
 cover/cover_bench_test.go | 78 +++++++++++++++++++++++++++++++++++++++
 cover/cover_test.go       | 15 ++++++++
 3 files changed, 100 insertions(+)
 create mode 100644 cover/cover_bench_test.go

diff --git a/cover/cover.go b/cover/cover.go
index be7ace4133..e512c8653f 100644
--- a/cover/cover.go
+++ b/cover/cover.go
@@ -31,6 +31,13 @@ func (c *Cover) Enabled() bool {
 	return true
 }
 
+// Config returns the standard Tracer configuration for the Cover tracer
+func (c *Cover) Config() topdown.TraceConfig {
+	return topdown.TraceConfig{
+		PlugLocalVars: false, // Event variable metadata is not required for the Coverage report
+	}
+}
+
 // Report returns a coverage Report for the given modules.
 func (c *Cover) Report(modules map[string]*ast.Module) (report Report) {
 	report.Files = map[string]*FileReport{}
diff --git a/cover/cover_bench_test.go b/cover/cover_bench_test.go
new file mode 100644
index 0000000000..dceb461345
--- /dev/null
+++ b/cover/cover_bench_test.go
@@ -0,0 +1,78 @@
+// Copyright 2020 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package cover
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/rego"
+)
+
+func BenchmarkCoverBigLocalVar(b *testing.B) {
+	iterations := []int{1, 100, 1000}
+	vars := []int{1, 10}
+
+	for _, iterationCount := range iterations {
+		for _, varCount := range vars {
+			name := fmt.Sprintf("%dVars%dIterations", varCount, iterationCount)
+			b.Run(name, func(b *testing.B) {
+				cover := New()
+				module := generateModule(varCount, iterationCount)
+
+				_, err := ast.ParseModule("test.rego", module)
+				if err != nil {
+					b.Fatal(err)
+				}
+
+				ctx := context.Background()
+
+				pq, err := rego.New(
+					rego.Module("test.rego", module),
+					rego.Query("data.test.p"),
+				).PrepareForEval(ctx)
+
+				if err != nil {
+					b.Fatal(err)
+				}
+
+				b.ResetTimer()
+
+				for i := 0; i < b.N; i++ {
+					b.StartTimer()
+					_, err = pq.Eval(ctx, rego.EvalTracer(cover))
+					b.StopTimer()
+
+					if err != nil {
+						b.Fatal(err)
+					}
+				}
+			})
+		}
+	}
+}
+
+func generateModule(numVars int, dataSize int) string {
+	sb := strings.Builder{}
+	sb.WriteString(`package test
+
+p {
+	x := a
+	v := x[i]
+`)
+	for i := 0; i < numVars; i++ {
+		sb.WriteString(fmt.Sprintf("\tv%d := x[i+%d]\n", i, i))
+	}
+	sb.WriteString("\tfalse\n}\n")
+	sb.WriteString("\na := [\n")
+	for i := 0; i < dataSize; i++ {
+		sb.WriteString(fmt.Sprintf("\t%d,\n", i))
+	}
+	sb.WriteString("]\n")
+	return sb.String()
+}
diff --git a/cover/cover_test.go b/cover/cover_test.go
index cd5df63ff3..97cc9ed24d 100644
--- a/cover/cover_test.go
+++ b/cover/cover_test.go
@@ -8,10 +8,12 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"reflect"
 	"testing"
 
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/rego"
+	"github.com/open-policy-agent/opa/topdown"
 )
 
 func TestCover(t *testing.T) {
@@ -136,3 +138,16 @@ p {
 		fmt.Println(string(bs))
 	}
 }
+
+func TestCoverTraceConfig(t *testing.T) {
+	ct := topdown.CustomTracer(New())
+	conf := ct.Config()
+
+	expected := topdown.TraceConfig{
+		PlugLocalVars: false,
+	}
+
+	if !reflect.DeepEqual(expected, conf) {
+		t.Fatalf("Expected config: %+v, got %+v", expected, conf)
+	}
+}