kala 2.1.0

use crate::{Duration, arch::GuardedArch, time};
use std::{
  hint::black_box,
  sync::{
    LazyLock,
    atomic::{Ordering, compiler_fence},
  },
};

// On a busy workstation that uses an Intel Core Ultra 7 265K, I've noticed zero
// difference in the computed overhead when varying OVERHEAD_SAMPLES to
// 100/1000/5000/10_000. None of it mattered. So we use small numbers to make it
// complete faster.
const OVERHEAD_SAMPLES: usize = 100;
// The Linux kernel needs 10-30 ms to ramp up the core's frequency to max. We
// conservatively use 50 ms.
const WARMUP_TIME: std::time::Duration = std::time::Duration::from_millis(50);

// The overhead of the measurement harness. Calculated by timing an empty
// closure.
//
// On an Intel Core Ultra 7 265K, I consistently observe an overhead of ~10
// nanoseconds.
//
// Note that the computed *overhead might be a zero-length Duration*. It can
// easily happen on CPUs with a very coarse TSC frequency.
static OVERHEAD: LazyLock<Duration> = LazyLock::new(|| {
  let stopwatch = Stopwatch::with_overhead(Duration::ZERO);
  let mut samples: Vec<Duration> = Vec::with_capacity(OVERHEAD_SAMPLES);

  let start_time = std::time::Instant::now();
  while start_time.elapsed() < WARMUP_TIME {
    for _ in 0..1000 {
      black_box(stopwatch.time(black_box(|| {})));
    }
  }

  for _ in 0..OVERHEAD_SAMPLES {
    samples.push(stopwatch.time(black_box(|| {})));
  }
  samples.sort(); // ascending

  // We use the Interquartile Mean (IQM) to estimate the true value since it is
  // both robust and stable (mean is not robust and median is not stable).
  let interquartile_start = (samples.len() as f64 * 0.25) as usize;
  let interquartile_end = (samples.len() as f64 * 0.75) as usize;
  let second_quartile = &samples[interquartile_start..interquartile_end];
  let mut second_quartile_sum = Duration::ZERO;
  for observation in second_quartile {
    second_quartile_sum = second_quartile_sum + *observation;
  }

  second_quartile_sum / second_quartile.len()
});

/// Allows precisely timing function executions.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub struct Stopwatch {
  arch: GuardedArch,
  overhead: Duration,
}

impl Stopwatch {
  // Computing `OVERHEAD` uses `Stopwatch`, but `Stopwatch::new()` uses
  // `OVERHEAD`. `Stopwatch::with_overhead` exists to break the cycle.
  fn with_overhead(overhead: Duration) -> Stopwatch {
    let sw = Stopwatch {
      arch: GuardedArch::best(),
      overhead,
    };

    // Trigger calibration if necessary on current architecture.
    let _ = sw.flushed_now();
    sw
  }

  /// Builds a new `Stopwatch`.
  ///
  /// # Implementation Details
  ///
  /// Automatic calibration will be performed on the first call to `new()` and
  /// then cached for the lifetime of the process. Depending on the CPU, this
  /// might take ~50-150 milliseconds.
  ///
  /// Subsequent calls to `new()` cost two `std::sync::LazyLock` (cached) reads.
  #[must_use]
  pub fn new() -> Stopwatch {
    Stopwatch::with_overhead(*OVERHEAD)
  }

  /// Returns true when a supported hardware time source is available.
  #[must_use]
  pub fn is_hw_supported(&self) -> bool {
    self.arch.is_hw_supported()
  }

  /// Returns the overhead of measuring an empty function.
  ///
  /// <div class="warning">
  ///
  /// This overhead is already subtracted out of the value returned by
  /// [`Stopwatch::time()`]. Do NOT subtract it out again!
  ///
  /// </div>
  #[must_use]
  pub fn overhead(&self) -> Duration {
    self.overhead
  }

  /// Runs the provided function _only once_ and returns the execution time.
  ///
  /// This might return a zero-length [`Duration`] IFF:
  ///
  /// - the provided function is very fast (few nanoseconds) AND,
  /// - the resolution of the best available clock is too coarse.
  ///
  /// Note that this function does NOT execute your function in a loop or
  /// estimate some sort of "average" function runtime. **The provided function
  /// is executed and timed just once**.
  ///
  /// The returned [`Duration`] already has [`Stopwatch::overhead()`] subtracted
  /// out.
  #[must_use]
  pub fn time<T, F: FnOnce() -> T>(&self, fun: F) -> Duration {
    let start = self.flushed_now();
    let retval = black_box(fun());
    let end = self.flushed_now();

    std::mem::drop(retval); // run `retval` Drop outside timed section
    // IMPORTANT: We subtract out the timing overhead (empty function "runtime")
    self.arch.elapsed(start, end) - self.overhead
  }

  // Get current time with compiler and clock (memory) fences to prevent
  // instruction and memory reordering.
  #[must_use]
  #[inline(always)]
  fn flushed_now(&self) -> time::Instant {
    self.arch.clock_fence();
    compiler_fence(Ordering::SeqCst);

    let instant = self.arch.now();

    compiler_fence(Ordering::SeqCst);
    self.arch.clock_fence();
    instant
  }
}

impl Default for Stopwatch {
  fn default() -> Self {
    Self::new()
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use crate::test_utils::assert_roughly_eq;
  use std::thread::sleep;
  use std::time::{Duration, Instant};

  const FUDGE_FACTOR: f64 = 0.05;

  #[test]
  fn basic() {
    let sw = Stopwatch::new();
    let mut std_nanos = 0.0;
    let our_nanos = sw
      .time(|| {
        let start = Instant::now();
        sleep(Duration::from_millis(10));
        std_nanos = start.elapsed().as_nanos() as f64;
      })
      .to_nanos();

    assert_roughly_eq(our_nanos, std_nanos, FUDGE_FACTOR);
  }
}