vortex_trace/

replay.rs

//! Generic deterministic replay framework.
//!
//! Since simulation is deterministic (same seed = same execution), we can
//! replay to any point and examine state. This module provides a generic
//! [`ReplayRunner`] that wraps any user-provided simulation function.
//!
//! Inspired by rr (Mozilla) — concept freely usable, independent implementation.

use crate::{SimTrace, TraceEvent};

/// A generic replay runner parameterized on user state type `S`.
///
/// The user provides a `sim_fn` that takes `(seed, target_tick)` and returns
/// `(state, trace)`. The replay runner uses this to implement time-travel
/// debugging primitives.
///
/// ```
/// use vortex_trace::replay::ReplayRunner;
/// use vortex_trace::SimTrace;
///
/// // Example: replay a simple counter simulation
/// let runner = ReplayRunner::new(42, |seed, ticks| {
///     let mut trace = SimTrace::new();
///     let mut counter = 0u64;
///     for t in 0..ticks {
///         counter += 1;
///         // Record events to trace as needed
///     }
///     (counter, trace)
/// });
///
/// let (state, _trace) = runner.replay_to_tick(100);
/// assert_eq!(state, 100);
/// ```
pub struct ReplayRunner<S, F>
where
    F: Fn(u64, u64) -> (S, SimTrace),
{
    seed: u64,
    sim_fn: F,
    _phantom: std::marker::PhantomData<S>,
}

impl<S, F> ReplayRunner<S, F>
where
    F: Fn(u64, u64) -> (S, SimTrace),
{
    /// Create a new replay runner.
    ///
    /// - `seed`: the deterministic seed for the simulation
    /// - `sim_fn`: `(seed, target_tick) -> (state, trace)`
    pub fn new(seed: u64, sim_fn: F) -> Self {
        Self {
            seed,
            sim_fn,
            _phantom: std::marker::PhantomData,
        }
    }

    /// Replay the simulation to a specific tick.
    ///
    /// Creates a fresh simulation from the seed and runs to `target_tick`.
    pub fn replay_to_tick(&self, target_tick: u64) -> (S, SimTrace) {
        (self.sim_fn)(self.seed, target_tick)
    }

    /// Replay a window: run to `end_tick`, return events in `[start_tick, end_tick]`.
    pub fn replay_window(&self, start_tick: u64, end_tick: u64) -> (S, Vec<TraceEvent>) {
        let (state, trace) = (self.sim_fn)(self.seed, end_tick);
        let window_events: Vec<TraceEvent> = trace
            .events()
            .iter()
            .filter(|e| e.tick >= start_tick && e.tick <= end_tick)
            .cloned()
            .collect();
        (state, window_events)
    }

    /// The seed used by this runner.
    pub fn seed(&self) -> u64 {
        self.seed
    }
}

/// Replay until a predicate is satisfied (binary-search friendly).
///
/// Runs the simulation at increasing tick counts until `predicate` returns
/// `true` or `max_ticks` is reached. Returns the tick where the predicate
/// was first satisfied.
///
/// This is more efficient than running tick-by-tick because it allows the
/// user's `sim_fn` to batch-run ticks internally.
pub fn replay_until<S, F, P>(
    seed: u64,
    max_ticks: u64,
    step: u64,
    sim_fn: F,
    predicate: P,
) -> Option<u64>
where
    F: Fn(u64, u64) -> (S, SimTrace),
    P: Fn(&S, &SimTrace) -> bool,
{
    let mut tick = step;
    while tick <= max_ticks {
        let (state, trace) = sim_fn(seed, tick);
        if predicate(&state, &trace) {
            // Found it — now binary search for the exact tick
            if tick <= step {
                return Some(tick);
            }
            let mut lo = tick - step;
            let mut hi = tick;
            while lo + 1 < hi {
                let mid = lo + (hi - lo) / 2;
                let (s, t) = sim_fn(seed, mid);
                if predicate(&s, &t) {
                    hi = mid;
                } else {
                    lo = mid;
                }
            }
            return Some(hi);
        }
        tick += step;
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::TraceEventKind;

    fn counter_sim(_seed: u64, ticks: u64) -> (u64, SimTrace) {
        let mut trace = SimTrace::new();
        let mut counter = 0u64;
        for t in 1..=ticks {
            counter += 1;
            if t % 10 == 0 {
                trace.record(
                    t,
                    1,
                    TraceEventKind::Custom {
                        tag: "counter".into(),
                        data: format!("{counter}"),
                    },
                );
            }
        }
        (counter, trace)
    }

    #[test]
    fn test_replay_to_tick() {
        let runner = ReplayRunner::new(42, counter_sim);
        let (state, _trace) = runner.replay_to_tick(100);
        assert_eq!(state, 100);
    }

    #[test]
    fn test_replay_deterministic() {
        let runner = ReplayRunner::new(42, counter_sim);
        let (s1, t1) = runner.replay_to_tick(500);
        let (s2, t2) = runner.replay_to_tick(500);
        assert_eq!(s1, s2);
        assert_eq!(t1.len(), t2.len());
    }

    #[test]
    fn test_replay_window() {
        let runner = ReplayRunner::new(42, counter_sim);
        let (state, events) = runner.replay_window(50, 100);
        assert_eq!(state, 100);
        // Events at ticks 50, 60, 70, 80, 90, 100 — 6 events in window
        assert_eq!(events.len(), 6);
        assert!(events.iter().all(|e| e.tick >= 50 && e.tick <= 100));
    }

    #[test]
    fn test_replay_until_found() {
        let tick = replay_until(42, 1000, 10, counter_sim, |state, _trace| *state >= 75);
        assert!(tick.is_some());
        assert_eq!(tick.unwrap(), 75);
    }

    #[test]
    fn test_replay_until_not_found() {
        let tick = replay_until(42, 50, 10, counter_sim, |state, _trace| *state >= 100);
        assert!(tick.is_none());
    }

    #[test]
    fn test_replay_seed() {
        let runner = ReplayRunner::new(123, counter_sim);
        assert_eq!(runner.seed(), 123);
    }
}