// Copyright 2016 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 topdown

import (
	"bytes"
	"fmt"
	"io"
	"maps"
	"slices"
	"strings"

	iStrs "github.com/open-policy-agent/opa/internal/strings"

	"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 (
	minLocationWidth      = len("query")
	maxIdealLocationWidth = 64
	columnPadding         = 4
	maxExprVarWidth       = 32
	maxPrettyExprVarWidth = 64
)

// Op defines the types of tracing events.
type Op string

const (
	// EnterOp is emitted when a new query is about to be evaluated.
	EnterOp Op = "Enter"

	// ExitOp is emitted when a query has evaluated to true.
	ExitOp Op = "Exit"

	// EvalOp is emitted when an expression is about to be evaluated.
	EvalOp Op = "Eval"

	// RedoOp is emitted when an expression, rule, or query is being re-evaluated.
	RedoOp Op = "Redo"

	// SaveOp is emitted when an expression is saved instead of evaluated
	// during partial evaluation.
	SaveOp Op = "Save"

	// FailOp is emitted when an expression evaluates to false.
	FailOp Op = "Fail"

	// DuplicateOp is emitted when a query has produced a duplicate value. The search
	// will stop at the point where the duplicate was emitted and backtrack.
	DuplicateOp Op = "Duplicate"

	// NoteOp is emitted when an expression invokes a tracing built-in function.
	NoteOp Op = "Note"

	// IndexOp is emitted during an expression evaluation to represent lookup
	// matches.
	IndexOp Op = "Index"

	// WasmOp is emitted when resolving a ref using an external
	// Resolver.
	WasmOp Op = "Wasm"

	// 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"
	FailedAssertionOp Op = "FailedAssertion"
)

// VarMetadata provides some user facing information about
// a variable in some policy.
type VarMetadata struct {
	Name     ast.Var       `json:"name"`
	Location *ast.Location `json:"location"`
}

// Event contains state associated with a tracing event.
type Event struct {
	Op            Op                      // Identifies type of event.
	Node          ast.Node                // Contains AST node relevant to the event.
	Location      *ast.Location           // The location of the Node this event relates to.
	QueryID       uint64                  // Identifies the query this event belongs to.
	ParentID      uint64                  // Identifies the parent query this event belongs to.
	Locals        *ast.ValueMap           // Contains local variable bindings from the query context. Nil if variables were not included in the trace event.
	LocalMetadata map[ast.Var]VarMetadata // Contains metadata for the local variable bindings. Nil if variables were not included in the trace event.
	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
	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.
func (evt *Event) HasRule() bool {
	_, ok := evt.Node.(*ast.Rule)
	return ok
}

// HasBody returns true if the Event contains an ast.Body.
func (evt *Event) HasBody() bool {
	_, ok := evt.Node.(ast.Body)
	return ok
}

// HasExpr returns true if the Event contains an ast.Expr.
func (evt *Event) HasExpr() bool {
	_, ok := evt.Node.(*ast.Expr)
	return ok
}

// Equal returns true if this event is equal to the other event.
func (evt *Event) Equal(other *Event) bool {
	if evt.Op != other.Op {
		return false
	}
	if evt.QueryID != other.QueryID {
		return false
	}
	if evt.ParentID != other.ParentID {
		return false
	}
	if !evt.equalNodes(other) {
		return false
	}
	return evt.Locals.Equal(other.Locals)
}

func (evt *Event) String() string {
	return fmt.Sprintf("%v %v %v (qid=%v, pqid=%v)", evt.Op, evt.Node, evt.Locals, evt.QueryID, evt.ParentID)
}

// Input returns the input object as it was at the event.
func (evt *Event) Input() *ast.Term {
	return evt.input
}

// Plug plugs event bindings into the provided ast.Term. Because bindings are mutable, this only makes sense to do when
// the event is emitted rather than on recorded trace events as the bindings are going to be different by then.
func (evt *Event) Plug(term *ast.Term) *ast.Term {
	return evt.bindings.Plug(term)
}

func (evt *Event) equalNodes(other *Event) bool {
	switch a := evt.Node.(type) {
	case ast.Body:
		if b, ok := other.Node.(ast.Body); ok {
			return a.Equal(b)
		}
	case *ast.Rule:
		if b, ok := other.Node.(*ast.Rule); ok {
			return a.Equal(b)
		}
	case *ast.Expr:
		if b, ok := other.Node.(*ast.Expr); ok {
			return a.Equal(b)
		}
	case nil:
		return other.Node == nil
	}
	return false
}

// Tracer defines the interface for tracing in the top-down evaluation engine.
//
// Deprecated: Use QueryTracer instead.
type Tracer interface {
	Enabled() bool
	Trace(*Event)
}

// QueryTracer defines the interface for tracing in the top-down evaluation engine.
// The implementation can provide additional configuration to modify the tracing
// behavior for query evaluations.
type QueryTracer interface {
	Enabled() bool
	TraceEvent(Event)
	Config() TraceConfig
}

// TraceConfig defines some common configuration for Tracer implementations
type TraceConfig struct {
	PlugLocalVars bool // Indicate whether to plug local variable bindings before calling into the tracer.
}

// legacyTracer Implements the QueryTracer interface by wrapping an older Tracer instance.
type legacyTracer struct {
	t Tracer
}

func (l *legacyTracer) Enabled() bool {
	return l.t.Enabled()
}

func (*legacyTracer) Config() TraceConfig {
	return TraceConfig{
		PlugLocalVars: true, // For backwards compatibility old tracers will plug local variables
	}
}

func (l *legacyTracer) TraceEvent(evt Event) {
	l.t.Trace(&evt)
}

// WrapLegacyTracer will create a new QueryTracer which wraps an
// older Tracer instance.
func WrapLegacyTracer(tracer Tracer) QueryTracer {
	return &legacyTracer{t: tracer}
}

// BufferTracer implements the Tracer and QueryTracer interface by
// simply buffering all events received.
type BufferTracer []*Event

// NewBufferTracer returns a new BufferTracer.
func NewBufferTracer() *BufferTracer {
	return &BufferTracer{}
}

// Enabled always returns true if the BufferTracer is instantiated.
func (b *BufferTracer) Enabled() bool {
	return b != nil
}

// Trace adds the event to the buffer.
//
// Deprecated: Use TraceEvent instead.
func (b *BufferTracer) Trace(evt *Event) {
	*b = append(*b, evt)
}

// TraceEvent adds the event to the buffer.
func (b *BufferTracer) TraceEvent(evt Event) {
	*b = append(*b, &evt)
}

// Config returns the Tracers standard configuration
func (*BufferTracer) Config() TraceConfig {
	return TraceConfig{PlugLocalVars: true}
}

// PrettyTrace pretty prints the trace to the writer.
func PrettyTrace(w io.Writer, trace []*Event) {
	PrettyTraceWithOpts(w, trace, PrettyTraceOptions{})
}

// PrettyTraceWithLocation prints the trace to the writer and includes location information
func PrettyTraceWithLocation(w io.Writer, trace []*Event) {
	PrettyTraceWithOpts(w, trace, PrettyTraceOptions{Locations: true})
}

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

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 {
				_, _ = w.Write(util.StringToByteSlice(cell))
			}
		}
		_, _ = fmt.Fprintln(w)
	}
}

func PrettyTraceWithOpts(w io.Writer, trace []*Event, opts PrettyTraceOptions) {
	depths := depths{}

	// FIXME: Can we shorten each location as we process each trace event instead of beforehand?
	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 {
			row = append(row, formatLocation(event, filePathAliases))
		}

		row = append(row, formatEvent(event, depth))

		if opts.ExprVariables {
			vars := exprLocalVars(event)
			keys := sortedKeys(vars)

			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))
				}
			}

			buf.WriteByte('}')
			row = append(row, buf.String())
		}

		if opts.LocalVariables {
			if locals := event.Locals; locals.Len() > 0 {
				keys := sortedKeys(locals)

				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.WriteByte('}')
				row = append(row, buf.String())
			} else {
				row = append(row, "{}")
			}
		}

		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
	})
	return util.SortedFunc(keys, func(a, b ast.Value) int {
		return strings.Compare(a.String(), b.String())
	})
}

func exprLocalVars(e *Event) *ast.ValueMap {
	vars := ast.NewValueMap()

	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.WalkVars(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.WalkVars(e.Node, findVars)

	return vars
}

func formatEvent(event *Event, depth int) string {
	buf := new(bytes.Buffer)
	formatEventPaddingAppend(buf, event, depth)
	buf.WriteString(string(event.Op))
	buf.WriteByte(' ')

	if event.Op == NoteOp {
		buf.WriteByte('"')
		buf.WriteString(event.Message)
		buf.WriteByte('"')

		return buf.String()
	}

	if node, ok := event.Node.(*ast.Rule); ok {
		bs, _ := node.Ref().ConstantPrefix().AppendText(buf.AvailableBuffer())
		buf.Write(bs)
	} else if event.Ref != nil {
		bs, _ := event.Ref.AppendText(buf.AvailableBuffer())
		buf.Write(bs)
	} else {
		fmt.Fprint(buf, rewrite(event).Node)
	}

	if event.Message != "" {
		buf.WriteByte(' ')
		buf.WriteString(event.Message)
	}

	return buf.String()
}

func formatEventPaddingAppend(buf *bytes.Buffer, event *Event, depth int) {
	for range formatEventSpaces(event, depth) - 1 {
		buf.WriteString("| ")
	}
}

func formatEventSpaces(event *Event, depth int) int {
	switch event.Op {
	case EnterOp:
		return depth
	case RedoOp:
		if _, ok := event.Node.(*ast.Expr); !ok {
			return depth
		}
	}
	return depth + 1
}

// getShortenedFileNames will return a map of file paths to shortened aliases
// that were found in the trace. It also returns the longest location expected
func getShortenedFileNames(trace []*Event) (map[string]string, int) {
	// Get a deduplicated list of all file paths
	// and the longest file path size
	fpAliases := map[string]string{}
	var canShorten []string
	longestLocation := 0
	for _, event := range trace {
		if event.Location != nil {
			if event.Location.File != "" {
				// length of "<name>:<row>"
				longestLocation = max(longestLocation, event.Location.StringLength())
				if _, ok := fpAliases[event.Location.File]; ok {
					continue
				}

				canShorten = append(canShorten, event.Location.File)

				// Default to just alias their full path
				fpAliases[event.Location.File] = event.Location.File
			} else {
				// length of "<min width>:<row>"
				longestLocation = max(longestLocation, minLocationWidth+util.NumDigitsInt(event.Location.Row)+1)
			}
		}
	}

	if len(canShorten) > 0 && longestLocation > maxIdealLocationWidth {
		fpAliases, longestLocation = iStrs.TruncateFilePaths(maxIdealLocationWidth, longestLocation, canShorten...)
	}

	return fpAliases, longestLocation
}

func formatLocation(event *Event, fileAliases map[string]string) string {
	if event.Location == nil {
		return ""
	}

	if event.Location.File == "" {
		return fmt.Sprintf("query:%v", event.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
// same query all have the same depth. The depth of query is
// depth(parent(query))+1.
type depths map[uint64]int

func (ds depths) GetOrSet(qid uint64, pqid uint64) int {
	depth := ds[qid]
	if depth == 0 {
		depth = ds[pqid]
		depth++
		ds[qid] = depth
	}
	return depth
}

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)
	}

	if !bctx.TraceEnabled {
		return iter(ast.InternedTerm(true))
	}

	evt := Event{
		Op:       NoteOp,
		Location: bctx.Location,
		QueryID:  bctx.QueryID,
		ParentID: bctx.ParentID,
		Message:  string(str),
	}

	for i := range bctx.QueryTracers {
		bctx.QueryTracers[i].TraceEvent(evt)
	}

	return iter(ast.InternedTerm(true))
}

func rewrite(event *Event) *Event {

	cpy := *event

	var node ast.Node

	switch v := event.Node.(type) {
	case *ast.Expr:
		expr := v.Copy()

		// Hide generated local vars in 'key' position that have not been
		// rewritten.
		if ev, ok := v.Terms.(*ast.Every); ok {
			if kv, ok := ev.Key.Value.(ast.Var); ok {
				if rw, ok := cpy.LocalMetadata[kv]; !ok || rw.Name.IsGenerated() {
					expr.Terms.(*ast.Every).Key = nil
				}
			}
		}
		node = expr
	case ast.Body:
		node = v.Copy()
	case *ast.Rule:
		node = v.Copy()
	}

	_, _ = ast.TransformVars(node, func(v ast.Var) (ast.Value, error) {
		if meta, ok := cpy.LocalMetadata[v]; ok {
			return meta.Name, nil
		}
		return v, nil
	})

	cpy.Node = node

	return &cpy
}

type varInfo struct {
	VarMetadata
	val     ast.Value
	exprLoc *ast.Location
	col     int // 0-indexed column
}

// resolveLocalRef resolves a ground ref whose base is a local variable (e.g. 'tc.data')
// against the given local bindings, returning the selected value. It returns nil if the
// ref can't be resolved: the base isn't a bound local, a path element isn't ground (after
// resolving any variable keys), or the path doesn't exist in the value.
func resolveLocalRef(ref ast.Ref, locals *ast.ValueMap) ast.Value {
	if len(ref) < 2 || locals == nil {
		return nil
	}

	base, ok := ref[0].Value.(ast.Var)
	if !ok {
		return nil
	}

	baseVal := locals.Get(base)
	if baseVal == nil {
		return nil
	}

	path := make(ast.Ref, 0, len(ref)-1)
	for _, t := range ref[1:] {
		if key, ok := t.Value.(ast.Var); ok {
			// A variable key (e.g. 'y[i]') must itself be resolved from the local bindings.
			keyVal := locals.Get(key)
			if keyVal == nil {
				return nil
			}
			path = append(path, ast.NewTerm(keyVal))
			continue
		}
		if !t.IsGround() {
			return nil
		}
		path = append(path, t)
	}

	val, err := baseVal.Find(path)
	if err != nil {
		return nil
	}
	return val
}

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 util.ByteSliceToString(v.exprLoc.Text)
	}
	return string(v.Name)
}

func padLocationText(loc *ast.Location) string {
	if loc == nil {
		return ""
	}

	text := util.ByteSliceToString(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 any, f func(*ast.Term) bool) {
	var vis *ast.GenericVisitor
	vis = ast.NewGenericVisitor(func(x any) 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.Ref:
				// For a ref that selects into a local variable (e.g. 'tc.data'), report the
				// selected value rather than only the base variable's (potentially large) value.
				// We keep descending (return false) so the base variable is still reported too.
				if val := resolveLocalRef(v, e.Locals); val != nil {
					info := varInfo{
						VarMetadata: VarMetadata{Name: ast.Var(term.Location.Text)},
						val:         val,
						exprLoc:     term.Location,
						col:         term.Location.Col,
					}
					if existing, exists := exprVars[info.Title()]; !exists || existing.val == nil {
						exprVars[info.Title()] = info
					}
				}
			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)
	w.Write(buf.Bytes())
	return nil
}

func printPrettyVars(w *bytes.Buffer, exprVars map[string]varInfo) {
	containsTabs := false
	varRows := make(map[int]any, len(exprVars))
	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.
		w.WriteString("\n\nWhere:\n")

		byName := slices.SortedStableFunc(maps.Values(exprVars), cmpVarInfoTitle)
		for _, info := range byName {
			fmt.Fprintf(w, "\n%s: %s", info.Title(), iStrs.Truncate(info.Value(), maxPrettyExprVarWidth))
		}

		return
	}

	byCol := util.SortedFunc(util.Values(exprVars), 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.WriteByte('\n')
	printArrows(w, byCol, -1)
	for i := range slices.Backward(byCol) {
		w.WriteByte('\n')
		printArrows(w, byCol, i)
	}
}

func cmpVarInfoTitle(a, b varInfo) int {
	return strings.Compare(a.Title(), b.Title())
}

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 := range spaces {
			var space byte = ' '
			if slices.Contains(info.exprLoc.Tabs, j+prevCol+1) {
				space = '\t'
			}
			w.WriteByte(space)
		}

		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(info.Title())
				w.WriteString(": ")
				w.WriteString(valueStr)
			} else {
				w.WriteString(valueStr)
			}
		} else {
			w.WriteByte('|')
		}
		prevCol = col
		isFirst = false
	}
}

func init() {
	RegisterBuiltinFunc(ast.Trace.Name, builtinTrace)
}