topdown: Fix trace to set location on notes

Previously the note events would not have a location on them which
mean they were difficult to track down (you would have to grep for the
message and hope it shows up.) With this change we just include the
location on notes like all other events.

Signed-off-by: Torin Sandall <torinsandall@gmail.com>

repl/repl_test.go

@@ -2032,10 +2032,10 @@ func TestEvalNotes(t *testing.T) {
 	repl.OneShot(ctx, "p")
 	expected := strings.TrimSpace(`query:1     Enter data.repl.p = _
 query:1     | Enter data.repl.p
-note        | | Note "x = 2"
+query:1     | | Note "x = 2"
 query:1     Redo data.repl.p = _
 query:1     | Redo data.repl.p
-note        | | Note "x = 3"
+query:1     | | Note "x = 3"
 true`)
 	expected += "\n"
 	if expected != buffer.String() {

topdown/trace.go

@@ -306,9 +306,6 @@ func numDigits10(n int) int {
 }
 
 func formatLocation(event *Event, fileAliases map[string]string) string {
-	if event.Op == NoteOp {
-		return fmt.Sprintf("%v", "note")
-	}
 
 	location := event.Location
 	if location == nil {
@@ -350,6 +347,7 @@ func builtinTrace(bctx BuiltinContext, args []*ast.Term, iter func(*ast.Term) er
 
 	evt := &Event{
 		Op:       NoteOp,
+		Location: bctx.Location,
 		QueryID:  bctx.QueryID,
 		ParentID: bctx.ParentID,
 		Message:  string(str),

topdown/trace_test.go

@@ -245,7 +245,7 @@ func TestPrettyTraceWithLocationTruncatedPaths(t *testing.T) {
 	compiler := ast.MustCompileModules(map[string]string{
 		"authz_bundle/com/foo/bar/baz/qux/acme/corp/internal/authz/policies/abac/v1/beta/policy.rego": `
 		package test
-		
+
 		import data.utils.q
 
 		p = true { q[x]; plus(x, 1, n) }
@@ -349,13 +349,13 @@ func TestPrettyTracePartialWithLocationTruncatedPaths(t *testing.T) {
 		package example_rbac
 
 		default allow = false
-		
+
 		allow {
 		    data.utils.user_has_role[role_name]
-		
+
 		    data.utils.role_has_permission[role_name]
 		}
-		
+
 		`,
 		"authz_bundle/com/foo/bar/baz/qux/acme/corp/internal/authz/policies/utils/user.rego": `
 		package utils
@@ -365,7 +365,7 @@ func TestPrettyTracePartialWithLocationTruncatedPaths(t *testing.T) {
 		    role_binding.role = role_name
 		    role_binding.user = input.subject.user
 		}
-		
+
 		role_has_permission[role_name] {
 		    role = data.roles[_]
 		    role.name = role_name
@@ -679,7 +679,7 @@ query:4     | | | Exit data.test.q
 query:3     | | Eval plus(x, 1, n)
 query:3     | | Eval sprintf("n= %v", [n], __local0__)
 query:3     | | Eval trace(__local0__)
-note        | | Note "n= 2"
+query:3     | | Note "n= 2"
 query:3     | | Exit data.test.p
 query:1     | Exit data.test.p = _
 query:1     Redo data.test.p = _
@@ -695,7 +695,7 @@ query:4     | | | Exit data.test.q
 query:3     | | Eval plus(x, 1, n)
 query:3     | | Eval sprintf("n= %v", [n], __local0__)
 query:3     | | Eval trace(__local0__)
-note        | | Note "n= 3"
+query:3     | | Note "n= 3"
 query:3     | | Exit data.test.p
 query:3     | Redo data.test.p
 query:3     | | Redo trace(__local0__)
@@ -708,7 +708,7 @@ query:4     | | | Exit data.test.q
 query:3     | | Eval plus(x, 1, n)
 query:3     | | Eval sprintf("n= %v", [n], __local0__)
 query:3     | | Eval trace(__local0__)
-note        | | Note "n= 4"
+query:3     | | Note "n= 4"
 query:3     | | Exit data.test.p
 query:3     | Redo data.test.p
 query:3     | | Redo trace(__local0__)
@@ -721,7 +721,7 @@ query:4     | | | Exit data.test.q
 query:3     | | Eval plus(x, 1, n)
 query:3     | | Eval sprintf("n= %v", [n], __local0__)
 query:3     | | Eval trace(__local0__)
-note        | | Note "n= 5"
+query:3     | | Note "n= 5"
 query:3     | | Exit data.test.p
 query:3     | Redo data.test.p
 query:3     | | Redo trace(__local0__)