A few tracing/profiling improvements

internal/presentation/presentation.go

@@ -35,8 +35,12 @@ import (
 	"github.com/open-policy-agent/opa/v1/topdown"
 )
 
-// DefaultProfileSortOrder is the default ordering unless something is specified in the CLI
-var DefaultProfileSortOrder = []string{"total_time_ns", "num_eval", "num_redo", "file", "line"}
+var (
+	// DefaultProfileSortOrder is the default ordering unless something is specified in the CLI
+	DefaultProfileSortOrder = []string{"total_time_ns", "num_eval", "num_redo", "file", "line"}
+
+	statKeys = []string{"min", "max", "mean", "90%", "99%"}
+)
 
 // DepAnalysisOutput contains the result of dependency analysis to be presented.
 type DepAnalysisOutput struct {
@@ -535,8 +539,6 @@ func prettyMetrics(w io.Writer, m metrics.Metrics, limit int) error {
 	return nil
 }
 
-var statKeys = []string{"min", "max", "mean", "90%", "99%"}
-
 func prettyAggregatedMetrics(w io.Writer, ms map[string]any, limit int) error {
 	keys := make([]string, 1, 1+len(statKeys))
 	keys[0] = "metric"

v1/profiler/profiler.go

@@ -6,6 +6,7 @@
 package profiler
 
 import (
+	"slices"
 	"sort"
 	"time"
 
@@ -14,6 +15,8 @@ import (
 	"github.com/open-policy-agent/opa/v1/topdown"
 )
 
+var unknownLocation = ast.NewLocation([]byte("???"), "", 0, 0)
+
 // Profiler computes and reports on the time spent on expressions.
 type Profiler struct {
 	hits            map[string]map[int]ExprStats
@@ -149,11 +152,7 @@ func (p *Profiler) Trace(event *topdown.Event) {
 // TraceEvent updates the coverage state.
 func (p *Profiler) TraceEvent(event topdown.Event) {
 	switch event.Op {
-	case topdown.EvalOp:
-		if expr, ok := event.Node.(*ast.Expr); ok && expr != nil {
-			p.processExpr(expr, event.Op)
-		}
-	case topdown.RedoOp:
+	case topdown.EvalOp, topdown.RedoOp:
 		if expr, ok := event.Node.(*ast.Expr); ok && expr != nil {
 			p.processExpr(expr, event.Op)
 		}
@@ -163,7 +162,7 @@ func (p *Profiler) TraceEvent(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)
+		expr.Location = unknownLocation
 	}
 
 	// set the active timer on the first expression
@@ -224,10 +223,12 @@ func (p *Profiler) processLastExpr() {
 func (p *Profiler) calculateHitsByExprIndex() {
 	file := p.prevExpr.location.File
 	hitsUnique, ok := p.hitsByExprIndex[file]
-
 	if !ok {
-		hitsUnique = map[int]map[int]ExprStats{}
-		hitsUnique[p.prevExpr.location.Row] = map[int]ExprStats{p.prevExpr.index: getProfilerStats(p.prevExpr, p.activeTimer)}
+		hitsUnique = map[int]map[int]ExprStats{
+			p.prevExpr.location.Row: {
+				p.prevExpr.index: getProfilerStats(p.prevExpr, p.activeTimer),
+			},
+		}
 		p.hitsByExprIndex[file] = hitsUnique
 	} else {
 		row := p.prevExpr.location.Row
@@ -323,8 +324,8 @@ func AggregateProfiles(profiles ...[]ExprStats) []ExprStatsAggregated {
 }
 
 func sortStatsByRow(ps []ExprStats) {
-	sort.Slice(ps, func(i, j int) bool {
-		return ps[i].Location.Row < ps[j].Location.Row
+	slices.SortFunc(ps, func(stat1, stat2 ExprStats) int {
+		return stat1.Location.Row - stat2.Location.Row
 	})
 }
 

v1/topdown/trace.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/open-policy-agent/opa/v1/ast"
 	"github.com/open-policy-agent/opa/v1/topdown/builtins"
+	"github.com/open-policy-agent/opa/v1/util"
 )
 
 const (
@@ -292,7 +293,7 @@ func (t *traceTable) write(w io.Writer, padding int) {
 			if i < len(row)-1 {
 				_, _ = fmt.Fprintf(w, "%-*s ", width, cell)
 			} else {
-				_, _ = fmt.Fprintf(w, "%s", cell)
+				_, _ = w.Write(util.StringToByteSlice(cell))
 			}
 		}
 		_, _ = fmt.Fprintln(w)
@@ -306,14 +307,14 @@ func PrettyTraceWithOpts(w io.Writer, trace []*Event, opts PrettyTraceOptions) {
 	filePathAliases, _ := getShortenedFileNames(trace)
 
 	table := traceTable{}
+	buf := new(bytes.Buffer)
 
 	for _, event := range trace {
 		depth := depths.GetOrSet(event.QueryID, event.ParentID)
 		row := traceRow{}
 
 		if opts.Locations {
-			location := formatLocation(event, filePathAliases)
-			row.add(location)
+			row.add(formatLocation(event, filePathAliases))
 		}
 
 		row.add(formatEvent(event, depth))
@@ -322,31 +323,47 @@ func PrettyTraceWithOpts(w io.Writer, trace []*Event, opts PrettyTraceOptions) {
 			vars := exprLocalVars(event)
 			keys := sortedKeys(vars)
 
-			buf := new(bytes.Buffer)
-			buf.WriteString("{")
-			for i, k := range keys {
-				if i > 0 {
+			buf.Reset()
+			buf.WriteByte('{')
+
+			if len(keys) > 0 {
+				k := keys[0]
+				buf.WriteString(k.String())
+				buf.WriteString(": ")
+				buf.WriteString(iStrs.Truncate(vars.Get(k).String(), maxExprVarWidth))
+
+				for _, k := range keys[1:] {
 					buf.WriteString(", ")
+					buf.WriteString(k.String())
+					buf.WriteString(": ")
+					buf.WriteString(iStrs.Truncate(vars.Get(k).String(), maxExprVarWidth))
 				}
-				_, _ = fmt.Fprintf(buf, "%v: %s", k, iStrs.Truncate(vars.Get(k).String(), maxExprVarWidth))
 			}
-			buf.WriteString("}")
+
+			buf.WriteByte('}')
 			row.add(buf.String())
 		}
 
 		if opts.LocalVariables {
-			if locals := event.Locals; locals != nil {
+			if locals := event.Locals; locals.Len() > 0 {
 				keys := sortedKeys(locals)
 
-				buf := new(bytes.Buffer)
-				buf.WriteString("{")
-				for i, k := range keys {
-					if i > 0 {
-						buf.WriteString(", ")
-					}
-					_, _ = fmt.Fprintf(buf, "%v: %s", k, iStrs.Truncate(locals.Get(k).String(), maxExprVarWidth))
+				buf.Reset()
+				buf.WriteByte('{')
+
+				k := keys[0]
+				buf.WriteString(k.String())
+				buf.WriteString(": ")
+				buf.WriteString(iStrs.Truncate(locals.Get(k).String(), maxExprVarWidth))
+
+				for _, k := range keys[1:] {
+					buf.WriteString(", ")
+					buf.WriteString(k.String())
+					buf.WriteString(": ")
+					buf.WriteString(iStrs.Truncate(locals.Get(k).String(), maxExprVarWidth))
 				}
-				buf.WriteString("}")
+
+				buf.WriteByte('}')
 				row.add(buf.String())
 			} else {
 				row.add("{}")
@@ -365,29 +382,26 @@ func sortedKeys(vm *ast.ValueMap) []ast.Value {
 		keys = append(keys, k)
 		return false
 	})
-	slices.SortFunc(keys, func(a, b ast.Value) int {
+	return util.SortedFunc(keys, func(a, b ast.Value) int {
 		return strings.Compare(a.String(), b.String())
 	})
-	return keys
 }
 
 func exprLocalVars(e *Event) *ast.ValueMap {
 	vars := ast.NewValueMap()
 
-	findVars := func(term *ast.Term) bool {
-		if name, ok := term.Value.(ast.Var); ok {
-			if meta, ok := e.LocalMetadata[name]; ok {
-				if val := e.Locals.Get(name); val != nil {
-					vars.Put(meta.Name, val)
-				}
+	findVars := func(name ast.Var) bool {
+		if meta, ok := e.LocalMetadata[name]; ok {
+			if val := e.Locals.Get(name); val != nil {
+				vars.Put(meta.Name, val)
 			}
 		}
 		return false
 	}
 
 	if r, ok := e.Node.(*ast.Rule); ok {
 		// We're only interested in vars in the head, not the body
-		ast.WalkTerms(r.Head, findVars)
+		ast.WalkVars(r.Head, findVars)
 		return vars
 	}
 
@@ -398,43 +412,47 @@ func exprLocalVars(e *Event) *ast.ValueMap {
 		return false
 	})
 
-	ast.WalkTerms(e.Node, findVars)
+	ast.WalkVars(e.Node, findVars)
 
 	return vars
 }
 
 func formatEvent(event *Event, depth int) string {
-	padding := formatEventPadding(event, depth)
+	buf := new(bytes.Buffer)
+	formatEventPaddingAppend(buf, event, depth)
+	buf.WriteString(string(event.Op))
+	buf.WriteByte(' ')
+
 	if event.Op == NoteOp {
-		return fmt.Sprintf("%v%v %q", padding, event.Op, event.Message)
+		buf.WriteByte('"')
+		buf.WriteString(event.Message)
+		buf.WriteByte('"')
+
+		return buf.String()
 	}
 
-	var details any
 	if node, ok := event.Node.(*ast.Rule); ok {
-		details = ast.RulePath(node)
+		bs, _ := node.Ref().ConstantPrefix().AppendText(buf.AvailableBuffer())
+		buf.Write(bs)
 	} else if event.Ref != nil {
-		details = event.Ref
+		bs, _ := event.Ref.AppendText(buf.AvailableBuffer())
+		buf.Write(bs)
 	} else {
-		details = rewrite(event).Node
+		fmt.Fprint(buf, rewrite(event).Node)
 	}
 
-	template := "%v%v %v"
-	opts := []any{padding, event.Op, details}
-
 	if event.Message != "" {
-		template += " %v"
-		opts = append(opts, event.Message)
+		buf.WriteByte(' ')
+		buf.WriteString(event.Message)
 	}
 
-	return fmt.Sprintf(template, opts...)
+	return buf.String()
 }
 
-func formatEventPadding(event *Event, depth int) string {
-	spaces := formatEventSpaces(event, depth)
-	if spaces > 1 {
-		return strings.Repeat("| ", spaces-1)
+func formatEventPaddingAppend(buf *bytes.Buffer, event *Event, depth int) {
+	for range formatEventSpaces(event, depth) - 1 {
+		buf.WriteString("| ")
 	}
-	return ""
 }
 
 func formatEventSpaces(event *Event, depth int) int {
@@ -461,11 +479,7 @@ func getShortenedFileNames(trace []*Event) (map[string]string, int) {
 		if event.Location != nil {
 			if event.Location.File != "" {
 				// length of "<name>:<row>"
-				curLen := len(event.Location.File) + numDigits10(event.Location.Row) + 1
-				if curLen > longestLocation {
-					longestLocation = curLen
-				}
-
+				longestLocation = max(longestLocation, event.Location.StringLength())
 				if _, ok := fpAliases[event.Location.File]; ok {
 					continue
 				}
@@ -476,10 +490,7 @@ func getShortenedFileNames(trace []*Event) (map[string]string, int) {
 				fpAliases[event.Location.File] = event.Location.File
 			} else {
 				// length of "<min width>:<row>"
-				curLen := minLocationWidth + numDigits10(event.Location.Row) + 1
-				if curLen > longestLocation {
-					longestLocation = curLen
-				}
+				longestLocation = max(longestLocation, minLocationWidth+util.NumDigitsInt(event.Location.Row)+1)
 			}
 		}
 	}
@@ -491,25 +502,16 @@ func getShortenedFileNames(trace []*Event) (map[string]string, int) {
 	return fpAliases, longestLocation
 }
 
-func numDigits10(n int) int {
-	if n < 10 {
-		return 1
-	}
-	return numDigits10(n/10) + 1
-}
-
 func formatLocation(event *Event, fileAliases map[string]string) string {
-
-	location := event.Location
-	if location == nil {
+	if event.Location == nil {
 		return ""
 	}
 
-	if location.File == "" {
-		return fmt.Sprintf("query:%v", location.Row)
+	if event.Location.File == "" {
+		return fmt.Sprintf("query:%v", event.Location.Row)
 	}
 
-	return fmt.Sprintf("%v:%v", fileAliases[location.File], location.Row)
+	return fmt.Sprintf("%v:%v", fileAliases[event.Location.File], event.Location.Row)
 }
 
 // depths is a helper for computing the depth of an event. Events within the
@@ -528,7 +530,6 @@ func (ds depths) GetOrSet(qid uint64, pqid uint64) int {
 }
 
 func builtinTrace(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
-
 	str, err := builtins.StringOperand(operands[0].Value, 1)
 	if err != nil {
 		return handleBuiltinErr(ast.Trace.Name, bctx.Location, err)