vortex_array/test_harness/trace/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Snapshot-friendly tracing harness for the optimizer and executor.
5//!
6//! # What this records
7//!
8//! [`trace_op`] runs a closure with a thread-local recorder installed. While the recorder is
9//! active, calls to the `trace_op!` macro inside the optimizer and executor push
10//! structured events into the recorder. The recorder produces a [`TraceDisplay`] that renders
11//! as a deterministic, hierarchical text trace suitable for `insta` snapshot assertions.
12//!
13//! Events cover:
14//!
15//! - **Optimization**: optimize/recursive-optimize entry, fixpoint loop iterations, applied
16//!   reduce rules, applied parent-reduce rules.
17//! - **Execution**: `execute_until` iterations, single-step entries, parent kernel attempts and
18//!   matches, slot transitions, builder start/append/finish, and the eventual canonical output.
19//!
20//! Despite the name `trace_op`, the harness is *not* a generic logging facility: it is closely
21//! coupled to the optimizer/executor state machines so that the resulting trace is stable enough
22//! to commit as a snapshot.
23//!
24//! # When to use it
25//!
26//! Use [`trace_op`] to write a regression test that asserts on the sequence of optimizer
27//! rewrites or executor steps an array goes through. Typical scenarios:
28//!
29//! - A reduce rule should fire exactly once on a specific input shape.
30//! - A parent kernel should be tried in a specific order and the first match should win.
31//! - The executor should walk into a slot, finish it, and pop back to the parent without
32//!   building a canonical intermediate.
33//! - A chunked array should drive the builder path rather than the stack path.
34//!
35//! Two resolutions are available:
36//!
37//! - [`TraceResolution::ExecutedOnly`] (default) — only events that actually fired (rule
38//!   rewrites that matched, kernels that succeeded, execution steps that ran). Optimizer
39//!   passes that produced no change are elided.
40//! - [`TraceResolution::Attempts`] — also records declined rule attempts, kernels that did
41//!   not match, and per-loop bookkeeping. Use this when ordering or fall-through matters.
42//!
43//! # Cost and scope
44//!
45//! - Capture is thread-local. Worker threads spawned inside `f` do not inherit the recorder.
46//! - Nested captures return an error so that unrelated traces never merge.
47//! - In release builds and CodSpeed benchmark builds, every `trace_op!` invocation is compiled
48//!   away by the macro's `cfg` gating; this module is then unused. See
49//!   `trace_op!` for the gating rules.
50//!
51//! # Example
52//!
53//! ```ignore
54//! use vortex_array::test_harness::trace::trace_op;
55//!
56//! let traced = trace_op(|| filter_array.optimize())?;
57//! assert!(traced.output.is::<Primitive>());
58//! insta::assert_snapshot!(traced.trace.to_string(), @r"
59//! optimize root=vortex.filter(i32, len=4) session=false
60//!   reduce TrivialFilterRule: vortex.filter(i32, len=4) -> vortex.primitive(i32, len=4)
61//!   done output=vortex.primitive(i32, len=4)
62//! ");
63//! ```
64
65use std::cell::Cell;
66use std::cell::RefCell;
67use std::fmt;
68use std::fmt::Debug;
69use std::fmt::Display;
70use std::sync::atomic::AtomicUsize;
71use std::sync::atomic::Ordering;
72
73use vortex_error::VortexExpect;
74use vortex_error::VortexResult;
75use vortex_error::vortex_err;
76
77use crate::ArrayRef;
78
79/// Controls how much rule and kernel resolution detail is captured.
80#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
81pub enum TraceResolution {
82    /// Record only the operations that actually executed.
83    #[default]
84    ExecutedOnly,
85    /// Also record rule and kernel attempts that matched but declined, or did not match.
86    Attempts,
87}
88
89/// Options for [`trace_op_with`].
90#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
91pub struct TraceOptions {
92    /// The amount of rule and kernel resolution detail to include.
93    pub resolution: TraceResolution,
94}
95
96/// The result of a traced operation.
97#[derive(Clone, Debug)]
98pub struct Traced<T> {
99    /// The value returned by the traced closure.
100    pub output: T,
101    /// A stable, snapshot-friendly rendering of optimizer and execution activity.
102    pub trace: TraceDisplay,
103}
104
105/// A stable, snapshot-friendly trace.
106#[derive(Clone, Debug, Default)]
107pub struct TraceDisplay {
108    options: TraceOptions,
109    events: Vec<TraceEvent>,
110}
111
112impl Display for TraceDisplay {
113    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
114        let hidden_events = self.hidden_events();
115        let mut optimize_depth = 0usize;
116        let mut wrote_event = false;
117
118        for (idx, event) in self.events.iter().enumerate() {
119            if hidden_events[idx] {
120                continue;
121            }
122
123            if event.closes_before(self.options.resolution) {
124                optimize_depth = optimize_depth.saturating_sub(1);
125            }
126
127            if event.is_hidden(self.options.resolution) {
128                continue;
129            }
130
131            if wrote_event {
132                writeln!(f)?;
133            } else {
134                wrote_event = true;
135            }
136
137            write_indent(
138                f,
139                optimize_depth + event.relative_indent(self.options.resolution, optimize_depth > 0),
140            )?;
141            event.fmt_line(f, self.options.resolution)?;
142
143            if event.opens_after(self.options.resolution) {
144                optimize_depth += 1;
145            }
146            if event.closes_after(self.options.resolution) {
147                optimize_depth = optimize_depth.saturating_sub(1);
148            }
149        }
150        Ok(())
151    }
152}
153
154impl TraceDisplay {
155    fn hidden_events(&self) -> Vec<bool> {
156        let mut hidden = vec![false; self.events.len()];
157        if self.options.resolution != TraceResolution::ExecutedOnly {
158            return hidden;
159        }
160
161        let mut optimize_stack = Vec::new();
162        for (idx, event) in self.events.iter().enumerate() {
163            match event {
164                TraceEvent::OptimizeStart { .. } => optimize_stack.push(idx),
165                TraceEvent::OptimizeDone { changed, .. } => {
166                    let Some(start) = optimize_stack.pop() else {
167                        continue;
168                    };
169                    if !changed {
170                        hidden[start..=idx].fill(true);
171                    }
172                }
173                _ => {}
174            }
175        }
176        hidden
177    }
178}
179
180fn write_indent(f: &mut fmt::Formatter<'_>, depth: usize) -> fmt::Result {
181    for _ in 0..depth {
182        f.write_str("  ")?;
183    }
184    Ok(())
185}
186
187/// Run `f` while capturing a trace of the optimizer and executor work it performs.
188///
189/// `f` typically invokes an operation that drives the executor or optimizer, such as
190/// [`ArrayOptimizer::optimize`][crate::optimizer::ArrayOptimizer::optimize] or
191/// [`ArrayRef::execute`][crate::ArrayRef::execute]. While `f` runs,
192/// the optimizer and executor emit structured events via the `trace_op!`
193/// macro into a thread-local recorder. When `f` returns, the recorder is finalized and
194/// returned alongside the closure's output as a [`Traced<T>`].
195///
196/// The default resolution ([`TraceResolution::ExecutedOnly`]) records the rule rewrites,
197/// parent kernels, execution steps, and builder activity that actually executed. Optimizer
198/// passes that produced no change are hidden from the rendered trace. Use [`trace_op_with`]
199/// with [`TraceResolution::Attempts`] when a test needs to assert on declined rule attempts,
200/// kernels that did not match, or other fall-through detail.
201///
202/// # Examples
203///
204/// ```ignore
205/// let traced = trace_op(|| filter_array.optimize())?;
206/// assert!(traced.output.is::<Primitive>());
207/// insta::assert_snapshot!(traced.trace.to_string(), @r"
208/// optimize root=vortex.filter(i32, len=4) session=false
209///   reduce TrivialFilterRule: vortex.filter(i32, len=4) -> vortex.primitive(i32, len=4)
210///   done output=vortex.primitive(i32, len=4)
211/// ");
212/// ```
213///
214/// # Errors
215///
216/// Returns whatever error `f` produces. Returns an error if a recorder is already active on
217/// the current thread — nested traces are not supported.
218pub fn trace_op<T>(f: impl FnOnce() -> VortexResult<T>) -> VortexResult<Traced<T>> {
219    trace_op_with(TraceOptions::default(), f)
220}
221
222/// Run `f` while capturing a trace using `options`.
223///
224/// See [`trace_op`] for the common case. Use this when you need to override the default
225/// [`TraceResolution`] to capture declined rules and unmatched kernels.
226///
227/// Trace capture is thread-local and intentionally does not propagate to worker threads. Nested
228/// trace captures return an error so tests do not accidentally merge unrelated traces.
229pub fn trace_op_with<T>(
230    options: TraceOptions,
231    f: impl FnOnce() -> VortexResult<T>,
232) -> VortexResult<Traced<T>> {
233    let interest = TraceInterest::from(options.resolution);
234    ACTIVE_TRACE.with(|active| {
235        let mut active = active.borrow_mut();
236        if active.is_some() {
237            return Err(vortex_err!("trace_op captures cannot be nested"));
238        }
239        *active = Some(TraceRecorder::new(options));
240        Ok(())
241    })?;
242    TRACE_INTEREST.with(|trace_interest| trace_interest.set(interest));
243    ACTIVE_TRACE_COUNT.fetch_add(1, Ordering::Relaxed);
244    if interest == TraceInterest::Attempts {
245        ATTEMPTS_TRACE_COUNT.fetch_add(1, Ordering::Relaxed);
246    }
247
248    let guard = ActiveTraceGuard { interest };
249    let output = f();
250    let recorder = ACTIVE_TRACE.with(|active| {
251        active
252            .borrow_mut()
253            .take()
254            .vortex_expect("trace recorder must be installed")
255    });
256    drop(guard);
257
258    output.map(|output| Traced {
259        output,
260        trace: recorder.finish(),
261    })
262}
263
264/// Returns true when the current thread has an active trace recorder.
265#[allow(clippy::inline_always)]
266#[inline(always)]
267pub(crate) fn is_active() -> bool {
268    if ACTIVE_TRACE_COUNT.load(Ordering::Relaxed) == 0 {
269        return false;
270    }
271    TRACE_INTEREST.with(|interest| interest.get().is_active())
272}
273
274#[allow(clippy::inline_always)]
275#[inline(always)]
276fn attempts_enabled() -> bool {
277    if ATTEMPTS_TRACE_COUNT.load(Ordering::Relaxed) == 0 {
278        return false;
279    }
280    TRACE_INTEREST.with(|interest| interest.get() == TraceInterest::Attempts)
281}
282
283#[derive(Clone, Copy, Debug)]
284enum TraceSource {
285    Static,
286    Session(usize),
287}
288
289impl Display for TraceSource {
290    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
291        match self {
292            TraceSource::Static => f.write_str("static"),
293            TraceSource::Session(idx) => write!(f, "session[{idx}]"),
294        }
295    }
296}
297
298#[derive(Clone, Copy, Debug)]
299enum AttemptOutcome {
300    Declined,
301    NoMatch,
302}
303
304impl Display for AttemptOutcome {
305    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
306        match self {
307            AttemptOutcome::Declined => f.write_str("declined"),
308            AttemptOutcome::NoMatch => f.write_str("no-match"),
309        }
310    }
311}
312
313#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
314enum TraceInterest {
315    #[default]
316    Off,
317    ExecutedOnly,
318    Attempts,
319}
320
321impl TraceInterest {
322    #[inline]
323    fn is_active(self) -> bool {
324        self != Self::Off
325    }
326}
327
328impl From<TraceResolution> for TraceInterest {
329    fn from(resolution: TraceResolution) -> Self {
330        match resolution {
331            TraceResolution::ExecutedOnly => Self::ExecutedOnly,
332            TraceResolution::Attempts => Self::Attempts,
333        }
334    }
335}
336
337/// Snapshot-friendly wrapper around [`ArrayRef`] that renders the encoding, dtype, and length
338/// using the canonical [`Display`] format (`vortex.primitive(i32, len=4)`).
339///
340/// Carries a clone of the [`ArrayRef`] instead of duplicating fields,
341/// so trace events stay small and share the same rendering as every other `{array}` print in
342/// the codebase.
343#[derive(Clone, Debug)]
344pub(crate) struct ArraySummary(ArrayRef);
345
346impl ArraySummary {
347    pub(crate) fn new(array: &ArrayRef) -> Self {
348        Self(array.clone())
349    }
350}
351
352impl Display for ArraySummary {
353    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
354        Display::fmt(&self.0, f)
355    }
356}
357
358pub(crate) fn record_optimize_start(root: &ArrayRef, session: bool) {
359    record(TraceEvent::OptimizeStart {
360        root: ArraySummary::new(root),
361        session,
362    });
363}
364
365pub(crate) fn record_optimize_loop_start(array: &ArrayRef) {
366    if !attempts_enabled() {
367        return;
368    }
369    record(TraceEvent::OptimizeLoopStart {
370        array: ArraySummary::new(array),
371    });
372}
373
374pub(crate) fn record_optimize_loop_end() {
375    if !attempts_enabled() {
376        return;
377    }
378    record(TraceEvent::OptimizeLoopEnd);
379}
380
381pub(crate) fn record_optimize_reduce_none(array: &ArrayRef) {
382    if !attempts_enabled() {
383        return;
384    }
385    record(TraceEvent::PhaseNone {
386        indent: 0,
387        phase: "reduce",
388        subject: "array",
389        array: ArraySummary::new(array),
390    });
391}
392
393pub(crate) fn record_optimize_parent_reduce_none(array: &ArrayRef) {
394    if !attempts_enabled() {
395        return;
396    }
397    record(TraceEvent::PhaseNone {
398        indent: 0,
399        phase: "reduce_parent",
400        subject: "array",
401        array: ArraySummary::new(array),
402    });
403}
404
405pub(crate) fn record_optimize_done(output: &ArrayRef, changed: bool) {
406    record(TraceEvent::OptimizeDone {
407        output: ArraySummary::new(output),
408        changed,
409    });
410}
411
412pub(crate) fn record_optimize_recursive_start(root: &ArrayRef) {
413    record(TraceEvent::OptimizeRecursiveStart {
414        root: ArraySummary::new(root),
415    });
416}
417
418pub(crate) fn record_optimize_recursive_slot(slot_idx: usize, input: &ArrayRef, output: &ArrayRef) {
419    record(TraceEvent::OptimizeRecursiveSlot {
420        slot_idx,
421        input: ArraySummary::new(input),
422        output: ArraySummary::new(output),
423    });
424}
425
426pub(crate) fn record_reduce_applied(array: &ArrayRef, rule: &dyn Debug, output: &ArrayRef) {
427    record(TraceEvent::ReduceApplied {
428        array: ArraySummary::new(array),
429        rule: compact_label(rule),
430        output: ArraySummary::new(output),
431    });
432}
433
434pub(crate) fn record_reduce_declined(array: &ArrayRef, rule: &dyn Debug) {
435    if !attempts_enabled() {
436        return;
437    }
438    record(TraceEvent::ReduceAttempt {
439        array: ArraySummary::new(array),
440        rule: compact_label(rule),
441        outcome: AttemptOutcome::Declined,
442    });
443}
444
445pub(crate) fn record_session_parent_reduce_applied(
446    parent: &ArrayRef,
447    child: &ArrayRef,
448    slot_idx: usize,
449    plugin_idx: usize,
450    output: &ArrayRef,
451) {
452    record_parent_reduce_applied(
453        parent,
454        child,
455        slot_idx,
456        TraceSource::Session(plugin_idx),
457        "reduce_parent_fn",
458        output,
459    );
460}
461
462pub(crate) fn record_session_parent_reduce_declined(
463    parent: &ArrayRef,
464    child: &ArrayRef,
465    slot_idx: usize,
466    plugin_idx: usize,
467) {
468    record_parent_reduce_attempt(
469        parent,
470        child,
471        slot_idx,
472        TraceSource::Session(plugin_idx),
473        "reduce_parent_fn",
474        AttemptOutcome::Declined,
475    );
476}
477
478pub(crate) fn record_static_parent_reduce_no_match(
479    parent: &ArrayRef,
480    child: &ArrayRef,
481    slot_idx: usize,
482    rule: &dyn Debug,
483) {
484    record_parent_reduce_attempt(
485        parent,
486        child,
487        slot_idx,
488        TraceSource::Static,
489        compact_label(rule),
490        AttemptOutcome::NoMatch,
491    );
492}
493
494pub(crate) fn record_static_parent_reduce_applied(
495    parent: &ArrayRef,
496    child: &ArrayRef,
497    slot_idx: usize,
498    rule: &dyn Debug,
499    output: &ArrayRef,
500) {
501    record_parent_reduce_applied(
502        parent,
503        child,
504        slot_idx,
505        TraceSource::Static,
506        compact_label(rule),
507        output,
508    );
509}
510
511pub(crate) fn record_static_parent_reduce_declined(
512    parent: &ArrayRef,
513    child: &ArrayRef,
514    slot_idx: usize,
515    rule: &dyn Debug,
516) {
517    record_parent_reduce_attempt(
518        parent,
519        child,
520        slot_idx,
521        TraceSource::Static,
522        compact_label(rule),
523        AttemptOutcome::Declined,
524    );
525}
526
527fn record_parent_reduce_attempt(
528    parent: &ArrayRef,
529    child: &ArrayRef,
530    slot_idx: usize,
531    source: TraceSource,
532    rule: impl Into<String>,
533    outcome: AttemptOutcome,
534) {
535    if !attempts_enabled() {
536        return;
537    }
538    record(TraceEvent::ParentReduceAttempt {
539        parent: ArraySummary::new(parent),
540        child: ArraySummary::new(child),
541        slot_idx,
542        source,
543        rule: rule.into(),
544        outcome,
545    });
546}
547
548fn record_parent_reduce_applied(
549    parent: &ArrayRef,
550    child: &ArrayRef,
551    slot_idx: usize,
552    source: TraceSource,
553    rule: impl Into<String>,
554    output: &ArrayRef,
555) {
556    record(TraceEvent::ParentReduceApplied {
557        parent: ArraySummary::new(parent),
558        child: ArraySummary::new(child),
559        slot_idx,
560        source,
561        rule: rule.into(),
562        output: ArraySummary::new(output),
563    });
564}
565
566pub(crate) fn record_execute_until_start<M>(root: &ArrayRef) {
567    record(TraceEvent::ExecuteUntilStart {
568        target: short_type_name::<M>(),
569        root: ArraySummary::new(root),
570    });
571}
572
573pub(crate) fn record_execute_until_iteration(
574    iteration: usize,
575    current: &ArrayRef,
576    stack_parent: Option<(&ArrayRef, usize)>,
577    builder_active: bool,
578) {
579    record(TraceEvent::ExecuteUntilIteration {
580        iteration,
581        current: ArraySummary::new(current),
582        stack_parent: stack_parent.map(|(array, slot_idx)| (ArraySummary::new(array), slot_idx)),
583        builder_active,
584    });
585}
586
587pub(crate) fn record_execute_until_done_check(target: bool, canonical: bool) {
588    if !attempts_enabled() {
589        return;
590    }
591    record(TraceEvent::ExecuteUntilDoneCheck { target, canonical });
592}
593
594pub(crate) fn record_execute_until_return(output: &ArrayRef) {
595    record(TraceEvent::ExecuteUntilReturn {
596        output: ArraySummary::new(output),
597    });
598}
599
600pub(crate) fn record_execute_until_pop_frame(slot_idx: usize, output: &ArrayRef) {
601    record(TraceEvent::ExecuteUntilPopFrame {
602        slot_idx,
603        output: ArraySummary::new(output),
604    });
605}
606
607pub(crate) fn record_session_execute_parent_applied(
608    phase: &'static str,
609    parent: &ArrayRef,
610    child: &ArrayRef,
611    slot_idx: usize,
612    plugin_idx: usize,
613    output: &ArrayRef,
614) {
615    record_execute_parent_applied(
616        phase,
617        parent,
618        child,
619        slot_idx,
620        TraceSource::Session(plugin_idx),
621        "execute_parent_fn",
622        output,
623    );
624}
625
626pub(crate) fn record_session_execute_parent_declined(
627    phase: &'static str,
628    parent: &ArrayRef,
629    child: &ArrayRef,
630    slot_idx: usize,
631    plugin_idx: usize,
632) {
633    record_execute_parent_attempt(
634        phase,
635        parent,
636        child,
637        slot_idx,
638        TraceSource::Session(plugin_idx),
639        "execute_parent_fn",
640        AttemptOutcome::Declined,
641    );
642}
643
644fn record_execute_parent_attempt(
645    phase: &'static str,
646    parent: &ArrayRef,
647    child: &ArrayRef,
648    slot_idx: usize,
649    source: TraceSource,
650    kernel: impl Into<String>,
651    outcome: AttemptOutcome,
652) {
653    if !attempts_enabled() {
654        return;
655    }
656    record(TraceEvent::ExecuteParentAttempt {
657        phase,
658        parent: ArraySummary::new(parent),
659        child: ArraySummary::new(child),
660        slot_idx,
661        source,
662        kernel: kernel.into(),
663        outcome,
664    });
665}
666
667fn record_execute_parent_applied(
668    phase: &'static str,
669    parent: &ArrayRef,
670    child: &ArrayRef,
671    slot_idx: usize,
672    source: TraceSource,
673    kernel: impl Into<String>,
674    output: &ArrayRef,
675) {
676    record(TraceEvent::ExecuteParentApplied {
677        phase,
678        parent: ArraySummary::new(parent),
679        child: ArraySummary::new(child),
680        slot_idx,
681        source,
682        kernel: kernel.into(),
683        output: ArraySummary::new(output),
684    });
685}
686
687pub(crate) fn record_execute_parent_none(phase: &'static str, current: &ArrayRef) {
688    if !attempts_enabled() {
689        return;
690    }
691    record(TraceEvent::PhaseNone {
692        indent: 2,
693        phase,
694        subject: "current",
695        array: ArraySummary::new(current),
696    });
697}
698
699pub(crate) fn record_execute_optimized(input: &ArrayRef, output: &ArrayRef) {
700    let changed = !ArrayRef::ptr_eq(input, output);
701    if !changed && !attempts_enabled() {
702        return;
703    }
704    record(TraceEvent::ExecuteOptimized {
705        input: ArraySummary::new(input),
706        output: ArraySummary::new(output),
707        changed,
708    });
709}
710
711pub(crate) fn record_execute_encoding(array: &ArrayRef) {
712    if !attempts_enabled() {
713        return;
714    }
715    record(TraceEvent::ExecuteEncoding {
716        array: ArraySummary::new(array),
717    });
718}
719
720pub(crate) fn record_execute_slot(slot_idx: usize, parent: &ArrayRef, child: &ArrayRef) {
721    record(TraceEvent::SlotTransition {
722        step: "ExecuteSlot",
723        slot_idx,
724        parent: ArraySummary::new(parent),
725        child: ArraySummary::new(child),
726    });
727}
728
729pub(crate) fn record_builder_start(array: &ArrayRef) {
730    record(TraceEvent::BuilderEvent {
731        action: "start",
732        subject: "array",
733        array: ArraySummary::new(array),
734    });
735}
736
737pub(crate) fn record_append_child(slot_idx: usize, parent: &ArrayRef, child: &ArrayRef) {
738    record(TraceEvent::SlotTransition {
739        step: "AppendChild",
740        slot_idx,
741        parent: ArraySummary::new(parent),
742        child: ArraySummary::new(child),
743    });
744}
745
746pub(crate) fn record_builder_append(child: &ArrayRef) {
747    record(TraceEvent::BuilderEvent {
748        action: "append",
749        subject: "child",
750        array: ArraySummary::new(child),
751    });
752}
753
754pub(crate) fn record_execute_done(array: &ArrayRef) {
755    record(TraceEvent::ExecuteDone {
756        array: ArraySummary::new(array),
757    });
758}
759
760pub(crate) fn record_builder_finish(output: &ArrayRef) {
761    record(TraceEvent::BuilderEvent {
762        action: "finish",
763        subject: "output",
764        array: ArraySummary::new(output),
765    });
766}
767
768pub(crate) fn record_single_step_start(array: &ArrayRef) {
769    record(TraceEvent::SingleStepStart {
770        array: ArraySummary::new(array),
771    });
772}
773
774pub(crate) fn record_single_step_phase_none(phase: &'static str, array: &ArrayRef) {
775    if !attempts_enabled() {
776        return;
777    }
778    record(TraceEvent::PhaseNone {
779        indent: 1,
780        phase,
781        subject: "array",
782        array: ArraySummary::new(array),
783    });
784}
785
786pub(crate) fn record_single_step_applied(phase: &'static str, input: &ArrayRef, output: &ArrayRef) {
787    record(TraceEvent::SingleStepApplied {
788        phase,
789        input: ArraySummary::new(input),
790        output: ArraySummary::new(output),
791    });
792}
793
794fn record(event: TraceEvent) {
795    ACTIVE_TRACE.with(|active| {
796        if let Some(recorder) = active.borrow_mut().as_mut() {
797            recorder.events.push(event);
798        }
799    });
800}
801
802fn compact_label(value: &dyn Debug) -> String {
803    let label = format!("{value:?}");
804    if let Some(label) = adapter_field(&label, "rule") {
805        return label.to_string();
806    }
807    if let Some(label) = adapter_field(&label, "kernel") {
808        return label.to_string();
809    }
810    label
811}
812
813fn adapter_field<'a>(label: &'a str, field: &str) -> Option<&'a str> {
814    let marker = format!("{field}: ");
815    let start = label.find(&marker)? + marker.len();
816    let rest = &label[start..];
817    let end = rest.rfind(" }")?;
818    Some(&rest[..end])
819}
820
821fn short_type_name<T>() -> String {
822    std::any::type_name::<T>()
823        .rsplit("::")
824        .next()
825        .vortex_expect("type names are never empty")
826        .to_string()
827}
828
829thread_local! {
830    static TRACE_INTEREST: Cell<TraceInterest> = const { Cell::new(TraceInterest::Off) };
831    static ACTIVE_TRACE: RefCell<Option<TraceRecorder>> = const { RefCell::new(None) };
832}
833
834static ACTIVE_TRACE_COUNT: AtomicUsize = AtomicUsize::new(0);
835static ATTEMPTS_TRACE_COUNT: AtomicUsize = AtomicUsize::new(0);
836
837struct ActiveTraceGuard {
838    interest: TraceInterest,
839}
840
841impl Drop for ActiveTraceGuard {
842    fn drop(&mut self) {
843        if self.interest == TraceInterest::Attempts {
844            ATTEMPTS_TRACE_COUNT.fetch_sub(1, Ordering::Relaxed);
845        }
846        ACTIVE_TRACE_COUNT.fetch_sub(1, Ordering::Relaxed);
847        TRACE_INTEREST.with(|interest| interest.set(TraceInterest::Off));
848        ACTIVE_TRACE.with(|active| {
849            active.borrow_mut().take();
850        });
851    }
852}
853
854#[derive(Debug)]
855struct TraceRecorder {
856    options: TraceOptions,
857    events: Vec<TraceEvent>,
858}
859
860impl TraceRecorder {
861    fn new(options: TraceOptions) -> Self {
862        Self {
863            options,
864            events: Vec::new(),
865        }
866    }
867
868    fn finish(self) -> TraceDisplay {
869        TraceDisplay {
870            options: self.options,
871            events: self.events,
872        }
873    }
874}
875
876#[derive(Clone, Debug)]
877enum TraceEvent {
878    OptimizeStart {
879        root: ArraySummary,
880        session: bool,
881    },
882    OptimizeLoopStart {
883        array: ArraySummary,
884    },
885    OptimizeLoopEnd,
886    OptimizeDone {
887        output: ArraySummary,
888        changed: bool,
889    },
890    OptimizeRecursiveStart {
891        root: ArraySummary,
892    },
893    OptimizeRecursiveSlot {
894        slot_idx: usize,
895        input: ArraySummary,
896        output: ArraySummary,
897    },
898    ReduceAttempt {
899        array: ArraySummary,
900        rule: String,
901        outcome: AttemptOutcome,
902    },
903    ReduceApplied {
904        array: ArraySummary,
905        rule: String,
906        output: ArraySummary,
907    },
908    ParentReduceAttempt {
909        parent: ArraySummary,
910        child: ArraySummary,
911        slot_idx: usize,
912        source: TraceSource,
913        rule: String,
914        outcome: AttemptOutcome,
915    },
916    ParentReduceApplied {
917        parent: ArraySummary,
918        child: ArraySummary,
919        slot_idx: usize,
920        source: TraceSource,
921        rule: String,
922        output: ArraySummary,
923    },
924    ExecuteUntilStart {
925        target: String,
926        root: ArraySummary,
927    },
928    ExecuteUntilIteration {
929        iteration: usize,
930        current: ArraySummary,
931        stack_parent: Option<(ArraySummary, usize)>,
932        builder_active: bool,
933    },
934    ExecuteUntilDoneCheck {
935        target: bool,
936        canonical: bool,
937    },
938    ExecuteUntilReturn {
939        output: ArraySummary,
940    },
941    ExecuteUntilPopFrame {
942        slot_idx: usize,
943        output: ArraySummary,
944    },
945    ExecuteParentAttempt {
946        phase: &'static str,
947        parent: ArraySummary,
948        child: ArraySummary,
949        slot_idx: usize,
950        source: TraceSource,
951        kernel: String,
952        outcome: AttemptOutcome,
953    },
954    ExecuteParentApplied {
955        phase: &'static str,
956        parent: ArraySummary,
957        child: ArraySummary,
958        slot_idx: usize,
959        source: TraceSource,
960        kernel: String,
961        output: ArraySummary,
962    },
963    PhaseNone {
964        indent: usize,
965        phase: &'static str,
966        subject: &'static str,
967        array: ArraySummary,
968    },
969    ExecuteOptimized {
970        input: ArraySummary,
971        output: ArraySummary,
972        changed: bool,
973    },
974    ExecuteEncoding {
975        array: ArraySummary,
976    },
977    SlotTransition {
978        step: &'static str,
979        slot_idx: usize,
980        parent: ArraySummary,
981        child: ArraySummary,
982    },
983    BuilderEvent {
984        action: &'static str,
985        subject: &'static str,
986        array: ArraySummary,
987    },
988    ExecuteDone {
989        array: ArraySummary,
990    },
991    SingleStepStart {
992        array: ArraySummary,
993    },
994    SingleStepApplied {
995        phase: &'static str,
996        input: ArraySummary,
997        output: ArraySummary,
998    },
999}
1000
1001impl TraceEvent {
1002    fn is_hidden(&self, resolution: TraceResolution) -> bool {
1003        match resolution {
1004            TraceResolution::Attempts => matches!(self, TraceEvent::OptimizeLoopEnd),
1005            TraceResolution::ExecutedOnly => matches!(
1006                self,
1007                TraceEvent::OptimizeLoopStart { .. }
1008                    | TraceEvent::OptimizeLoopEnd
1009                    | TraceEvent::PhaseNone { .. }
1010                    | TraceEvent::ExecuteUntilDoneCheck { .. }
1011                    | TraceEvent::ExecuteEncoding { .. }
1012                    | TraceEvent::ExecuteOptimized { changed: false, .. }
1013                    | TraceEvent::ExecuteParentAttempt { .. }
1014                    | TraceEvent::ReduceAttempt { .. }
1015                    | TraceEvent::ParentReduceAttempt { .. }
1016            ),
1017        }
1018    }
1019
1020    fn opens_after(&self, resolution: TraceResolution) -> bool {
1021        match resolution {
1022            TraceResolution::Attempts => matches!(
1023                self,
1024                TraceEvent::OptimizeStart { .. } | TraceEvent::OptimizeLoopStart { .. }
1025            ),
1026            TraceResolution::ExecutedOnly => matches!(self, TraceEvent::OptimizeStart { .. }),
1027        }
1028    }
1029
1030    fn closes_before(&self, resolution: TraceResolution) -> bool {
1031        match resolution {
1032            TraceResolution::Attempts => matches!(self, TraceEvent::OptimizeLoopEnd),
1033            TraceResolution::ExecutedOnly => false,
1034        }
1035    }
1036
1037    fn closes_after(&self, _resolution: TraceResolution) -> bool {
1038        matches!(self, TraceEvent::OptimizeDone { .. })
1039    }
1040
1041    fn relative_indent(&self, _resolution: TraceResolution, in_optimize_scope: bool) -> usize {
1042        match self {
1043            TraceEvent::OptimizeStart { .. }
1044            | TraceEvent::OptimizeLoopStart { .. }
1045            | TraceEvent::OptimizeDone { .. } => 0,
1046            TraceEvent::ReduceAttempt { .. }
1047            | TraceEvent::ReduceApplied { .. }
1048            | TraceEvent::ParentReduceAttempt { .. }
1049            | TraceEvent::ParentReduceApplied { .. }
1050                if in_optimize_scope =>
1051            {
1052                0
1053            }
1054            TraceEvent::PhaseNone { indent, .. } => *indent,
1055            TraceEvent::ReduceAttempt { .. }
1056            | TraceEvent::ReduceApplied { .. }
1057            | TraceEvent::ParentReduceAttempt { .. }
1058            | TraceEvent::ParentReduceApplied { .. }
1059            | TraceEvent::ExecuteUntilDoneCheck { .. }
1060            | TraceEvent::ExecuteUntilPopFrame { .. }
1061            | TraceEvent::ExecuteParentAttempt { .. }
1062            | TraceEvent::ExecuteParentApplied { .. }
1063            | TraceEvent::ExecuteOptimized { .. }
1064            | TraceEvent::ExecuteEncoding { .. }
1065            | TraceEvent::SlotTransition { .. }
1066            | TraceEvent::BuilderEvent { .. }
1067            | TraceEvent::ExecuteDone { .. } => 2,
1068            TraceEvent::OptimizeRecursiveSlot { .. }
1069            | TraceEvent::ExecuteUntilIteration { .. }
1070            | TraceEvent::ExecuteUntilReturn { .. }
1071            | TraceEvent::SingleStepApplied { .. } => 1,
1072            TraceEvent::OptimizeLoopEnd
1073            | TraceEvent::OptimizeRecursiveStart { .. }
1074            | TraceEvent::ExecuteUntilStart { .. }
1075            | TraceEvent::SingleStepStart { .. } => 0,
1076        }
1077    }
1078
1079    fn fmt_line(&self, f: &mut fmt::Formatter<'_>, resolution: TraceResolution) -> fmt::Result {
1080        match self {
1081            TraceEvent::OptimizeStart { root, session } => {
1082                write!(f, "optimize root={root} session={session}")
1083            }
1084            TraceEvent::OptimizeLoopStart { array } => {
1085                write!(f, "loop input={array}")
1086            }
1087            TraceEvent::OptimizeLoopEnd => Ok(()),
1088            TraceEvent::OptimizeDone { output, changed } => match resolution {
1089                TraceResolution::Attempts => write!(f, "done output={output} changed={changed}"),
1090                TraceResolution::ExecutedOnly => write!(f, "done output={output}"),
1091            },
1092            TraceEvent::OptimizeRecursiveStart { root } => {
1093                write!(f, "optimize_recursive root={root}")
1094            }
1095            TraceEvent::OptimizeRecursiveSlot {
1096                slot_idx,
1097                input,
1098                output,
1099            } => write!(f, "recursive slot={slot_idx} input={input} output={output}"),
1100            TraceEvent::ReduceAttempt {
1101                array,
1102                rule,
1103                outcome,
1104            } => write!(
1105                f,
1106                "reduce attempt array={array} source=static rule={rule} outcome={outcome}"
1107            ),
1108            TraceEvent::ReduceApplied {
1109                array,
1110                rule,
1111                output,
1112            } => match resolution {
1113                TraceResolution::Attempts => write!(
1114                    f,
1115                    "reduce applied array={array} source=static rule={rule} output={output}"
1116                ),
1117                TraceResolution::ExecutedOnly => {
1118                    write!(f, "reduce {rule}: {array} -> {output}")
1119                }
1120            },
1121            TraceEvent::ParentReduceAttempt {
1122                parent,
1123                child,
1124                slot_idx,
1125                source,
1126                rule,
1127                outcome,
1128            } => write!(
1129                f,
1130                "reduce_parent attempt slot={slot_idx} parent={parent} child={child} source={source} rule={rule} outcome={outcome}"
1131            ),
1132            TraceEvent::ParentReduceApplied {
1133                parent,
1134                child,
1135                slot_idx,
1136                source,
1137                rule,
1138                output,
1139            } => match resolution {
1140                TraceResolution::Attempts => write!(
1141                    f,
1142                    "reduce_parent applied slot={slot_idx} parent={parent} child={child} source={source} rule={rule} output={output}"
1143                ),
1144                TraceResolution::ExecutedOnly => write!(
1145                    f,
1146                    "reduce_parent {source}:{rule} slot={slot_idx} parent={parent} child={child} -> {output}"
1147                ),
1148            },
1149            TraceEvent::ExecuteUntilStart { target, root } => {
1150                write!(f, "execute_until target={target} root={root}")
1151            }
1152            TraceEvent::ExecuteUntilIteration {
1153                iteration,
1154                current,
1155                stack_parent,
1156                builder_active,
1157            } => {
1158                write!(f, "iter {iteration} current={current}")?;
1159                if let Some((parent, slot_idx)) = stack_parent {
1160                    write!(f, " stack_parent={parent} slot={slot_idx}")?;
1161                }
1162                write!(f, " builder_active={builder_active}")
1163            }
1164            TraceEvent::ExecuteUntilDoneCheck { target, canonical } => {
1165                write!(f, "done_check target={target} canonical={canonical}")
1166            }
1167            TraceEvent::ExecuteUntilReturn { output } => {
1168                write!(f, "return output={output}")
1169            }
1170            TraceEvent::ExecuteUntilPopFrame { slot_idx, output } => {
1171                write!(f, "pop_frame slot={slot_idx} output={output}")
1172            }
1173            TraceEvent::ExecuteParentAttempt {
1174                phase,
1175                parent,
1176                child,
1177                slot_idx,
1178                source,
1179                kernel,
1180                outcome,
1181            } => write!(
1182                f,
1183                "{phase} attempt slot={slot_idx} parent={parent} child={child} source={source} kernel={kernel} outcome={outcome}"
1184            ),
1185            TraceEvent::ExecuteParentApplied {
1186                phase,
1187                parent,
1188                child,
1189                slot_idx,
1190                source,
1191                kernel,
1192                output,
1193            } => match resolution {
1194                TraceResolution::Attempts => write!(
1195                    f,
1196                    "{phase} applied slot={slot_idx} parent={parent} child={child} source={source} kernel={kernel} output={output}"
1197                ),
1198                TraceResolution::ExecutedOnly => write!(
1199                    f,
1200                    "{phase} {source}:{kernel} slot={slot_idx} parent={parent} child={child} -> {output}"
1201                ),
1202            },
1203            TraceEvent::PhaseNone {
1204                phase,
1205                subject,
1206                array,
1207                ..
1208            } => {
1209                write!(f, "{phase} none {subject}={array}")
1210            }
1211            TraceEvent::ExecuteOptimized {
1212                input,
1213                output,
1214                changed,
1215            } => match resolution {
1216                TraceResolution::Attempts => write!(
1217                    f,
1218                    "optimize_ctx input={input} output={output} changed={changed}"
1219                ),
1220                TraceResolution::ExecutedOnly => write!(f, "optimize_ctx {input} -> {output}"),
1221            },
1222            TraceEvent::ExecuteEncoding { array } => {
1223                write!(f, "execute encoding={array}")
1224            }
1225            TraceEvent::SlotTransition {
1226                step,
1227                slot_idx,
1228                parent,
1229                child,
1230            } => write!(f, "{step} slot={slot_idx} parent={parent} child={child}"),
1231            TraceEvent::ExecuteDone { array } => {
1232                write!(f, "Done array={array}")
1233            }
1234            TraceEvent::BuilderEvent {
1235                action,
1236                subject,
1237                array,
1238            } => {
1239                write!(f, "builder {action} {subject}={array}")
1240            }
1241            TraceEvent::SingleStepStart { array } => {
1242                write!(f, "execute_step input={array}")
1243            }
1244            TraceEvent::SingleStepApplied {
1245                phase,
1246                input,
1247                output,
1248            } => match resolution {
1249                TraceResolution::Attempts => {
1250                    write!(f, "{phase} applied input={input} output={output}")
1251                }
1252                TraceResolution::ExecutedOnly => write!(f, "{phase} {input} -> {output}"),
1253            },
1254        }
1255    }
1256}
1257
1258#[cfg(test)]
1259mod tests;
vortex_array/test_harness/trace/mod.rs

vortex_array/test_harness/trace/
mod.rs
