update dependencies (#6267)

Signed-off-by: hongming <coder.scala@gmail.com>

vendor/github.com/open-policy-agent/opa/topdown/trace.go

@@ -5,8 +5,10 @@
 package topdown
 
 import (
+	"bytes"
 	"fmt"
 	"io"
+	"slices"
 	"strings"
 
 	iStrs "github.com/open-policy-agent/opa/internal/strings"
@@ -18,7 +20,9 @@ import (
 const (
 	minLocationWidth      = 5 // len("query")
 	maxIdealLocationWidth = 64
-	locationPadding       = 4
+	columnPadding         = 4
+	maxExprVarWidth       = 32
+	maxPrettyExprVarWidth = 64
 )
 
 // Op defines the types of tracing events.
@@ -62,7 +66,8 @@ const (
 	// UnifyOp is emitted when two terms are unified.  Node will be set to an
 	// equality expression with the two terms.  This Node will not have location
 	// info.
-	UnifyOp Op = "Unify"
+	UnifyOp           Op = "Unify"
+	FailedAssertionOp Op = "FailedAssertion"
 )
 
 // VarMetadata provides some user facing information about
@@ -84,8 +89,14 @@ type Event struct {
 	Message       string                  // Contains message for Note events.
 	Ref           *ast.Ref                // Identifies the subject ref for the event. Only applies to Index and Wasm operations.
 
-	input    *ast.Term
-	bindings *bindings
+	input                     *ast.Term
+	bindings                  *bindings
+	localVirtualCacheSnapshot *ast.ValueMap
+}
+
+func (evt *Event) WithInput(input *ast.Term) *Event {
+	evt.input = input
+	return evt
 }
 
 // HasRule returns true if the Event contains an ast.Rule.
@@ -236,31 +247,162 @@ func (b *BufferTracer) Config() TraceConfig {
 
 // PrettyTrace pretty prints the trace to the writer.
 func PrettyTrace(w io.Writer, trace []*Event) {
-	prettyTraceWith(w, trace, false)
+	PrettyTraceWithOpts(w, trace, PrettyTraceOptions{})
 }
 
 // PrettyTraceWithLocation prints the trace to the writer and includes location information
 func PrettyTraceWithLocation(w io.Writer, trace []*Event) {
-	prettyTraceWith(w, trace, true)
+	PrettyTraceWithOpts(w, trace, PrettyTraceOptions{Locations: true})
 }
 
-func prettyTraceWith(w io.Writer, trace []*Event, locations bool) {
+type PrettyTraceOptions struct {
+	Locations      bool // Include location information
+	ExprVariables  bool // Include variables found in the expression
+	LocalVariables bool // Include all local variables
+}
+
+type traceRow []string
+
+func (r *traceRow) add(s string) {
+	*r = append(*r, s)
+}
+
+type traceTable struct {
+	rows      []traceRow
+	maxWidths []int
+}
+
+func (t *traceTable) add(row traceRow) {
+	t.rows = append(t.rows, row)
+	for i := range row {
+		if i >= len(t.maxWidths) {
+			t.maxWidths = append(t.maxWidths, len(row[i]))
+		} else if len(row[i]) > t.maxWidths[i] {
+			t.maxWidths[i] = len(row[i])
+		}
+	}
+}
+
+func (t *traceTable) write(w io.Writer, padding int) {
+	for _, row := range t.rows {
+		for i, cell := range row {
+			width := t.maxWidths[i] + padding
+			if i < len(row)-1 {
+				_, _ = fmt.Fprintf(w, "%-*s ", width, cell)
+			} else {
+				_, _ = fmt.Fprintf(w, "%s", cell)
+			}
+		}
+		_, _ = fmt.Fprintln(w)
+	}
+}
+
+func PrettyTraceWithOpts(w io.Writer, trace []*Event, opts PrettyTraceOptions) {
 	depths := depths{}
 
-	filePathAliases, longest := getShortenedFileNames(trace)
+	// FIXME: Can we shorten each location as we process each trace event instead of beforehand?
+	filePathAliases, _ := getShortenedFileNames(trace)
 
-	// Always include some padding between the trace and location
-	locationWidth := longest + locationPadding
+	table := traceTable{}
 
 	for _, event := range trace {
 		depth := depths.GetOrSet(event.QueryID, event.ParentID)
-		if locations {
+		row := traceRow{}
+
+		if opts.Locations {
 			location := formatLocation(event, filePathAliases)
-			fmt.Fprintf(w, "%-*s %s\n", locationWidth, location, formatEvent(event, depth))
-		} else {
-			fmt.Fprintln(w, formatEvent(event, depth))
+			row.add(location)
 		}
+
+		row.add(formatEvent(event, depth))
+
+		if opts.ExprVariables {
+			vars := exprLocalVars(event)
+			keys := sortedKeys(vars)
+
+			buf := new(bytes.Buffer)
+			buf.WriteString("{")
+			for i, k := range keys {
+				if i > 0 {
+					buf.WriteString(", ")
+				}
+				_, _ = fmt.Fprintf(buf, "%v: %s", k, iStrs.Truncate(vars.Get(k).String(), maxExprVarWidth))
+			}
+			buf.WriteString("}")
+			row.add(buf.String())
+		}
+
+		if opts.LocalVariables {
+			if locals := event.Locals; locals != nil {
+				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.WriteString("}")
+				row.add(buf.String())
+			} else {
+				row.add("{}")
+			}
+		}
+
+		table.add(row)
 	}
+
+	table.write(w, columnPadding)
+}
+
+func sortedKeys(vm *ast.ValueMap) []ast.Value {
+	keys := make([]ast.Value, 0, vm.Len())
+	vm.Iter(func(k, _ ast.Value) bool {
+		keys = append(keys, k)
+		return false
+	})
+	slices.SortFunc(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 r, ok := term.Value.(ast.Ref); ok {
+		//	fmt.Printf("ref: %v\n", r)
+		//	//return true
+		//}
+		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)
+				}
+			}
+		}
+		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)
+		return vars
+	}
+
+	// The local cache snapshot only contains a snapshot for those refs present in the event node,
+	// so they can all be added to the vars map.
+	e.localVirtualCacheSnapshot.Iter(func(k, v ast.Value) bool {
+		vars.Put(k, v)
+		return false
+	})
+
+	ast.WalkTerms(e.Node, findVars)
+
+	return vars
 }
 
 func formatEvent(event *Event, depth int) string {
@@ -451,6 +593,310 @@ func rewrite(event *Event) *Event {
 	return &cpy
 }
 
+type varInfo struct {
+	VarMetadata
+	val     ast.Value
+	exprLoc *ast.Location
+	col     int // 0-indexed column
+}
+
+func (v varInfo) Value() string {
+	if v.val != nil {
+		return v.val.String()
+	}
+	return "undefined"
+}
+
+func (v varInfo) Title() string {
+	if v.exprLoc != nil && v.exprLoc.Text != nil {
+		return string(v.exprLoc.Text)
+	}
+	return string(v.Name)
+}
+
+func padLocationText(loc *ast.Location) string {
+	if loc == nil {
+		return ""
+	}
+
+	text := string(loc.Text)
+
+	if loc.Col == 0 {
+		return text
+	}
+
+	buf := new(bytes.Buffer)
+	j := 0
+	for i := 1; i < loc.Col; i++ {
+		if len(loc.Tabs) > 0 && j < len(loc.Tabs) && loc.Tabs[j] == i {
+			buf.WriteString("\t")
+			j++
+		} else {
+			buf.WriteString(" ")
+		}
+	}
+
+	buf.WriteString(text)
+	return buf.String()
+}
+
+type PrettyEventOpts struct {
+	PrettyVars bool
+}
+
+func walkTestTerms(x interface{}, f func(*ast.Term) bool) {
+	var vis *ast.GenericVisitor
+	vis = ast.NewGenericVisitor(func(x interface{}) bool {
+		switch x := x.(type) {
+		case ast.Call:
+			for _, t := range x[1:] {
+				vis.Walk(t)
+			}
+			return true
+		case *ast.Expr:
+			if x.IsCall() {
+				for _, o := range x.Operands() {
+					vis.Walk(o)
+				}
+				for i := range x.With {
+					vis.Walk(x.With[i])
+				}
+				return true
+			}
+		case *ast.Term:
+			return f(x)
+		case *ast.With:
+			vis.Walk(x.Value)
+			return true
+		}
+		return false
+	})
+	vis.Walk(x)
+}
+
+func PrettyEvent(w io.Writer, e *Event, opts PrettyEventOpts) error {
+	if !opts.PrettyVars {
+		_, _ = fmt.Fprintln(w, padLocationText(e.Location))
+		return nil
+	}
+
+	buf := new(bytes.Buffer)
+	exprVars := map[string]varInfo{}
+
+	findVars := func(unknownAreUndefined bool) func(term *ast.Term) bool {
+		return func(term *ast.Term) bool {
+			if term.Location == nil {
+				return false
+			}
+
+			switch v := term.Value.(type) {
+			case *ast.ArrayComprehension, *ast.SetComprehension, *ast.ObjectComprehension:
+				// we don't report on the internals of a comprehension, as it's already evaluated, and we won't have the local vars.
+				return true
+			case ast.Var:
+				var info *varInfo
+				if meta, ok := e.LocalMetadata[v]; ok {
+					info = &varInfo{
+						VarMetadata: meta,
+						val:         e.Locals.Get(v),
+						exprLoc:     term.Location,
+					}
+				} else if unknownAreUndefined {
+					info = &varInfo{
+						VarMetadata: VarMetadata{Name: v},
+						exprLoc:     term.Location,
+						col:         term.Location.Col,
+					}
+				}
+
+				if info != nil {
+					if v, exists := exprVars[info.Title()]; !exists || v.val == nil {
+						if term.Location != nil {
+							info.col = term.Location.Col
+						}
+						exprVars[info.Title()] = *info
+					}
+				}
+			}
+			return false
+		}
+	}
+
+	expr, ok := e.Node.(*ast.Expr)
+	if !ok || expr == nil {
+		return nil
+	}
+
+	base := expr.BaseCogeneratedExpr()
+	exprText := padLocationText(base.Location)
+	buf.WriteString(exprText)
+
+	e.localVirtualCacheSnapshot.Iter(func(k, v ast.Value) bool {
+		var info *varInfo
+		switch k := k.(type) {
+		case ast.Ref:
+			info = &varInfo{
+				VarMetadata: VarMetadata{Name: ast.Var(k.String())},
+				val:         v,
+				exprLoc:     k[0].Location,
+				col:         k[0].Location.Col,
+			}
+		case *ast.ArrayComprehension:
+			info = &varInfo{
+				VarMetadata: VarMetadata{Name: ast.Var(k.String())},
+				val:         v,
+				exprLoc:     k.Term.Location,
+				col:         k.Term.Location.Col,
+			}
+		case *ast.SetComprehension:
+			info = &varInfo{
+				VarMetadata: VarMetadata{Name: ast.Var(k.String())},
+				val:         v,
+				exprLoc:     k.Term.Location,
+				col:         k.Term.Location.Col,
+			}
+		case *ast.ObjectComprehension:
+			info = &varInfo{
+				VarMetadata: VarMetadata{Name: ast.Var(k.String())},
+				val:         v,
+				exprLoc:     k.Key.Location,
+				col:         k.Key.Location.Col,
+			}
+		}
+
+		if info != nil {
+			exprVars[info.Title()] = *info
+		}
+
+		return false
+	})
+
+	// If the expression is negated, we can't confidently assert that vars with unknown values are 'undefined',
+	// since the compiler might have opted out of the necessary rewrite.
+	walkTestTerms(expr, findVars(!expr.Negated))
+	coExprs := expr.CogeneratedExprs()
+	for _, coExpr := range coExprs {
+		// Only the current "co-expr" can have undefined vars, if we don't know the value for a var in any other co-expr,
+		// it's unknown, not undefined. A var can be unknown if it hasn't been assigned a value yet, because the co-expr
+		// hasn't been evaluated yet (the fail happened before it).
+		walkTestTerms(coExpr, findVars(false))
+	}
+
+	printPrettyVars(buf, exprVars)
+	_, _ = fmt.Fprint(w, buf.String())
+	return nil
+}
+
+func printPrettyVars(w *bytes.Buffer, exprVars map[string]varInfo) {
+	containsTabs := false
+	varRows := make(map[int]interface{})
+	for _, info := range exprVars {
+		if len(info.exprLoc.Tabs) > 0 {
+			containsTabs = true
+		}
+		varRows[info.exprLoc.Row] = nil
+	}
+
+	if containsTabs && len(varRows) > 1 {
+		// We can't (currently) reliably point to var locations when they are on different rows that contain tabs.
+		// So we'll just print them in alphabetical order instead.
+		byName := make([]varInfo, 0, len(exprVars))
+		for _, info := range exprVars {
+			byName = append(byName, info)
+		}
+		slices.SortStableFunc(byName, func(a, b varInfo) int {
+			return strings.Compare(a.Title(), b.Title())
+		})
+
+		w.WriteString("\n\nWhere:\n")
+		for _, info := range byName {
+			w.WriteString(fmt.Sprintf("\n%s: %s", info.Title(), iStrs.Truncate(info.Value(), maxPrettyExprVarWidth)))
+		}
+
+		return
+	}
+
+	byCol := make([]varInfo, 0, len(exprVars))
+	for _, info := range exprVars {
+		byCol = append(byCol, info)
+	}
+	slices.SortFunc(byCol, func(a, b varInfo) int {
+		// sort first by column, then by reverse row (to present vars in the same order they appear in the expr)
+		if a.col == b.col {
+			if a.exprLoc.Row == b.exprLoc.Row {
+				return strings.Compare(a.Title(), b.Title())
+			}
+			return b.exprLoc.Row - a.exprLoc.Row
+		}
+		return a.col - b.col
+	})
+
+	if len(byCol) == 0 {
+		return
+	}
+
+	w.WriteString("\n")
+	printArrows(w, byCol, -1)
+	for i := len(byCol) - 1; i >= 0; i-- {
+		w.WriteString("\n")
+		printArrows(w, byCol, i)
+	}
+}
+
+func printArrows(w *bytes.Buffer, l []varInfo, printValueAt int) {
+	prevCol := 0
+	var slice []varInfo
+	if printValueAt >= 0 {
+		slice = l[:printValueAt+1]
+	} else {
+		slice = l
+	}
+	isFirst := true
+	for i, info := range slice {
+
+		isLast := i >= len(slice)-1
+		col := info.col
+
+		if !isLast && col == l[i+1].col {
+			// We're sharing the same column with another, subsequent var
+			continue
+		}
+
+		spaces := col - 1
+		if i > 0 && !isFirst {
+			spaces = (col - prevCol) - 1
+		}
+
+		for j := 0; j < spaces; j++ {
+			tab := false
+			for _, t := range info.exprLoc.Tabs {
+				if t == j+prevCol+1 {
+					w.WriteString("\t")
+					tab = true
+					break
+				}
+			}
+			if !tab {
+				w.WriteString(" ")
+			}
+		}
+
+		if isLast && printValueAt >= 0 {
+			valueStr := iStrs.Truncate(info.Value(), maxPrettyExprVarWidth)
+			if (i > 0 && col == l[i-1].col) || (i < len(l)-1 && col == l[i+1].col) {
+				// There is another var on this column, so we need to include the name to differentiate them.
+				w.WriteString(fmt.Sprintf("%s: %s", info.Title(), valueStr))
+			} else {
+				w.WriteString(valueStr)
+			}
+		} else {
+			w.WriteString("|")
+		}
+		prevCol = col
+		isFirst = false
+	}
+}
+
 func init() {
 	RegisterBuiltinFunc(ast.Trace.Name, builtinTrace)
 }