From 4e5f7916734a17603e78e3464fd6b4e1ee6f9b59 Mon Sep 17 00:00:00 2001
From: Ashutosh Narkar <anarkar4387@gmail.com>
Date: Tue, 7 Apr 2020 16:47:45 -0700
Subject: [PATCH] profiler: Group expressions with missing locations

Previously the profiler would panic if it encountered an expression with no location info. This change groups such expressions by giving them a fake location so that their evaluation results can be captured.

Signed-off-by: Ashutosh Narkar <anarkar4387@gmail.com>
---
 profiler/profiler.go      |  4 +++
 profiler/profiler_test.go | 69 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/profiler/profiler.go b/profiler/profiler.go
index 66f644ae6e..836d0b5abb 100644
--- a/profiler/profiler.go
+++ b/profiler/profiler.go
@@ -130,6 +130,10 @@ func (p *Profiler) Trace(event *topdown.Event) {
 }
 
 func (p *Profiler) processExpr(expr *ast.Expr, eventType topdown.Op) {
+	if expr.Location == nil {
+		// add fake location to group expressions without a location
+		expr.Location = ast.NewLocation([]byte("???"), "", 0, 0)
+	}
 
 	// set the active timer on the first expression
 	if p.activeTimer.IsZero() {
diff --git a/profiler/profiler_test.go b/profiler/profiler_test.go
index 2aa994c6cb..378ae015cc 100644
--- a/profiler/profiler_test.go
+++ b/profiler/profiler_test.go
@@ -410,3 +410,72 @@ baz {
 		}
 	}
 }
+
+func TestProfilerWithPartialEval(t *testing.T) {
+	profiler := New()
+
+	module := `package test
+
+default foo = false
+
+foo = true {
+	op = allowed_operations[_]
+	input.method = op.method
+	input.resource = op.resource
+}
+
+allowed_operations = [
+	{"method": "PUT", "resource": "policy"},
+]`
+
+	_, err := ast.ParseModule("test.rego", module)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ctx := context.Background()
+
+	pq, err := rego.New(
+		rego.Module("test.rego", module),
+		rego.Query("data.test.foo"),
+	).PrepareForEval(ctx, rego.WithPartialEval())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = pq.Eval(ctx, rego.EvalTracer(profiler))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	report := profiler.ReportByFile()
+
+	if len(report.Files) != 1 {
+		t.Fatalf("Expected file report length to be 1 instead got %v", len(report.Files))
+	}
+
+	fr := report.Files[""]
+
+	if len(fr.Result) != 2 {
+		t.Fatalf("Expected 2 results for file but instead got %v", len(fr.Result))
+	}
+
+	expectedNumEval := []int{2, 1}
+	expectedNumRedo := []int{2, 1}
+	expectedLocation := []string{"???", "data.partial.__result__"}
+
+	for idx, actualExprStat := range fr.Result {
+		if actualExprStat.NumEval != expectedNumEval[idx] {
+			t.Fatalf("Index %v: Expected number of evals %v but got %v", idx, expectedNumEval[idx], actualExprStat.NumEval)
+		}
+
+		if actualExprStat.NumRedo != expectedNumRedo[idx] {
+			t.Fatalf("Index %v: Expected number of redos %v but got %v", idx, expectedNumRedo[idx], actualExprStat.NumRedo)
+		}
+
+		if string(actualExprStat.Location.Text) != expectedLocation[idx] {
+			t.Fatalf("Index %v: Expected location %v but got %v", idx, expectedLocation[idx], string(actualExprStat.Location.Text))
+		}
+
+	}
+}