kala 2.0.0

use std::cmp::Ordering;
use std::ops::{Add, Div, Sub};

#[derive(Debug, Clone, Copy)]
enum DurationImpl {
  Zero, // Useful because it's unitless.
  Ticks { ticks: u64, frequency: u64 },
  // NOTE: 2^64 nanos is ~585 years. It won't overflow when computing a
  // function's runtime.
  Nanos(u64),
}

impl Ord for DurationImpl {
  // NOTE: This implementation works for all `DurationImpl`s returned from
  // `Stopwatch`. Since we let the user create their own `Duration`s (and thus
  // `DurationImpl`s), this can produce weird results when comparing between
  // DurationImpls with different storage types. Don't mix those.
  fn cmp(&self, other: &Self) -> Ordering {
    fn normalize(value: DurationImpl) -> f64 {
      match value {
        DurationImpl::Zero => 0.0,
        DurationImpl::Ticks { ticks, frequency } => {
          ticks_to_nanos(ticks, frequency)
        },
        DurationImpl::Nanos(nanos) => nanos as f64,
      }
    }

    let self_normalized = normalize(*self);
    let other_normalized = normalize(*other);
    self_normalized.total_cmp(&other_normalized)
  }
}

impl PartialOrd for DurationImpl {
  fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
    Some(self.cmp(other))
  }
}

impl PartialEq for DurationImpl {
  fn eq(&self, other: &Self) -> bool {
    self.cmp(other) == Ordering::Equal
  }
}

impl Eq for DurationImpl {} // Marker trait.

/// Describes the possible internal storage types for a [`Duration`].
#[derive(Debug, PartialEq, Eq, Clone, Copy, PartialOrd, Ord, Hash)]
pub enum StorageType {
  /// A zero-length `Duration`.
  Zero,
  /// Stores ticks and frequency.
  Ticks,
  /// Stores nanoseconds.
  Nanos,
}

/// Represents a span of time.
///
/// Similar in principle to [`std::time::Duration`], but supports different
/// internal representations that are more granular (and precise) when hardware
/// clock counters like the _[Time Stamp Counter][tsc]_ (TSC) on [x86-64] or the
/// _[Cycle Counter Register][ccnt]_ (CCNT) on [AArch64] are available.
///
/// This type implements [`Add`], [`Sub`], and [`Div`], supporting basic
/// addition, subtraction and division. **Subtraction saturates on underflow**,
/// unlike [`std::time::Duration`]. Negative `Duration`s are not possible.
/// Division by zero, negative numbers or similar silliness will panic.
///
/// You may wish to keep intermediate values as this type as long as possible
/// before converting to another unit to avoid conversion/rounding losses. E.g.
/// if calculating an average function runtime: sum up all of the collected
/// `Duration`s, divide the sum by the count of `Duration`s and only call
/// [`Duration::to_nanos()`] at the very end. You might need [`Duration::ZERO`].
///
/// This type implements [`Ord`], [`PartialEq`] and related comparison traits.
/// The implementations of the comparison traits will work perfectly fine for
/// all `Duration`s produced by [`Stopwatch`][crate::Stopwatch], but care should
/// be taken when comparing user-created `Duration`s (those returned from
/// [`Duration::from_nanos()`] and [`Duration::from_ticks()`]). In general, do
/// NOT mix `Duration`s with different storage types.
///
/// [x86-64]: https://en.wikipedia.org/wiki/X86-64
/// [aarch64]: https://en.wikipedia.org/wiki/AArch64
/// [tsc]: https://en.wikipedia.org/wiki/Time_Stamp_Counter
/// [ccnt]: https://support.arm.com/documentation/ddi0211/k/system-control-coprocessor/system-control-coprocessor-register-descriptions/c15--cycle-counter-register--ccnt-
#[derive(Debug, PartialEq, Eq, Clone, Copy, PartialOrd, Ord)]
pub struct Duration {
  inner: DurationImpl,
}

impl Duration {
  /// A zero-length span of time.
  ///
  /// Useful because it is internally _unitless_. It can be added to or
  /// subtracted from another `Duration` _without_ triggering conversions.
  pub const ZERO: Duration = Duration {
    inner: DurationImpl::Zero,
  };

  /// Create a `Duration` from a number of nanoseconds.
  ///
  /// If `nanos` is 0, then [`Duration::ZERO`] is returned.
  ///
  /// <div class="warning">
  ///
  /// Do NOT mix `Duration`s created with [`Duration::from_nanos()`] and those
  /// created with [`Duration::from_ticks()`]. It can lead to weird results from
  /// comparison traits ([`Ord`], [`PartialEq`] and similar).
  ///
  /// User-created `Duration`s are only supported to make it easier for you to
  /// write tests that use `Duration`s.
  ///
  /// </div>
  #[must_use]
  pub fn from_nanos(nanos: u64) -> Duration {
    match nanos {
      0 => Duration::ZERO,
      _ => Duration {
        inner: DurationImpl::Nanos(nanos),
      },
    }
  }

  /// Create a `Duration` from a number of clock ticks of a specific frequency.
  ///
  /// If `ticks` is 0, then [`Duration::ZERO`] is returned.
  ///
  /// <div class="warning">
  ///
  /// Do NOT mix `Duration`s created with [`Duration::from_nanos()`] and those
  /// created with [`Duration::from_ticks()`]. It can lead to weird results from
  /// comparison traits ([`Ord`], [`PartialEq`] and similar).
  ///
  /// User-created `Duration`s are only supported to make it easier for you to
  /// write tests that use `Duration`s.
  ///
  /// </div>
  #[must_use]
  pub fn from_ticks(ticks: u64, frequency: u64) -> Duration {
    match ticks {
      0 => Duration::ZERO,
      _ => Duration {
        inner: DurationImpl::Ticks { ticks, frequency },
      },
    }
  }

  /// Elapsed time in _possibly fractional_ nanoseconds if the underlying data
  /// support the precision.
  #[must_use]
  pub fn to_nanos(&self) -> f64 {
    match self.inner {
      DurationImpl::Nanos(nanos) => nanos as f64,
      DurationImpl::Ticks { ticks, frequency } => {
        ticks_to_nanos(ticks, frequency)
      },
      DurationImpl::Zero => 0.0,
    }
  }

  /// Returns the internal storage type used for this `Duration`.
  #[must_use]
  pub fn storage_type(&self) -> StorageType {
    match self.inner {
      DurationImpl::Zero => StorageType::Zero,
      DurationImpl::Ticks { .. } => StorageType::Ticks,
      DurationImpl::Nanos(_) => StorageType::Nanos,
    }
  }

  /// Elapsed time in *opaque time units* that are only fit for
  /// relative comparison ("this time delta is smaller/bigger than this other
  /// delta").
  ///
  /// The returned value is the _most granular unit_ that could be computed
  /// _without any conversion losses_.
  ///
  /// # Implementation details
  ///
  /// The returned value may be in nanoseconds, in _invariant_ CPU clock
  /// ticks, or something else entirely.
  ///
  /// <div class="warning">
  ///
  /// _No guarantees are provided_ about the returned values. The magnitude of
  /// the values could change from one release to the next _including in new
  /// `PATCH` versions!_
  ///
  /// Do NOT store the values in any kind of storage that will outlive the OS
  /// process.
  ///
  /// </div>
  #[must_use]
  pub fn to_raw(&self) -> u64 {
    match self.inner {
      DurationImpl::Ticks { ticks, .. } => ticks,
      DurationImpl::Nanos(nanos) => nanos,
      DurationImpl::Zero => 0,
    }
  }

  /// Elapsed time converted to an [`std::time::Duration`].
  ///
  /// If the underlying data are more precise than nanoseconds, the value is
  /// rounded to the nearest nanosecond (since [`std::time::Duration`] does not
  /// support values smaller than whole nanoseconds).
  ///
  /// See [`Duration::to_nanos()`] for access to fractional nanoseconds.
  #[must_use]
  pub fn to_std_duration(&self) -> std::time::Duration {
    let nanos = match self.inner {
      DurationImpl::Nanos(nanos) => nanos,
      DurationImpl::Ticks { ticks, frequency } => {
        ticks_to_nanos(ticks, frequency).round() as u64
      },
      DurationImpl::Zero => 0,
    };
    std::time::Duration::from_nanos(nanos)
  }
}

impl Add for Duration {
  type Output = Duration;

  fn add(self, rhs: Self) -> Self::Output {
    match (self.inner, rhs.inner) {
      (DurationImpl::Nanos(self_nanos), DurationImpl::Nanos(rhs_nanos)) => {
        Duration::from_nanos(self_nanos + rhs_nanos)
      },
      (
        DurationImpl::Ticks { ticks, frequency },
        DurationImpl::Ticks {
          ticks: rhs_ticks, ..
        },
      ) => Duration::from_ticks(ticks + rhs_ticks, frequency),
      (DurationImpl::Zero, other) => Duration { inner: other },
      (other, DurationImpl::Zero) => Duration { inner: other },
      (
        DurationImpl::Nanos(nanos),
        DurationImpl::Ticks { ticks, frequency },
      ) => Duration::from_nanos(
        nanos + ticks_to_nanos(ticks, frequency).round() as u64,
      ),
      (
        DurationImpl::Ticks { ticks, frequency },
        DurationImpl::Nanos(nanos),
      ) => Duration::from_nanos(
        ticks_to_nanos(ticks, frequency).round() as u64 + nanos,
      ),
    }
  }
}

impl Sub for Duration {
  type Output = Duration;

  fn sub(self, rhs: Self) -> Self::Output {
    match (self.inner, rhs.inner) {
      (DurationImpl::Nanos(self_nanos), DurationImpl::Nanos(rhs_nanos)) => {
        Duration::from_nanos(self_nanos.saturating_sub(rhs_nanos))
      },
      (
        DurationImpl::Ticks { ticks, frequency },
        DurationImpl::Ticks {
          ticks: rhs_ticks, ..
        },
      ) => Duration::from_ticks(ticks.saturating_sub(rhs_ticks), frequency),
      (DurationImpl::Zero, _) => Duration::ZERO,
      (other, DurationImpl::Zero) => Duration { inner: other },
      (
        DurationImpl::Nanos(nanos),
        DurationImpl::Ticks { ticks, frequency },
      ) => Duration::from_nanos(
        nanos - ticks_to_nanos(ticks, frequency).round() as u64,
      ),
      (
        DurationImpl::Ticks { ticks, frequency },
        DurationImpl::Nanos(nanos),
      ) => Duration::from_nanos(
        ticks_to_nanos(ticks, frequency).round() as u64 - nanos,
      ),
    }
  }
}

impl<T> Div<T> for Duration
where
  T: TryInto<u64> + Copy,
  <T as std::convert::TryInto<u64>>::Error: std::fmt::Debug,
{
  type Output = Duration;

  fn div(self, rhs: T) -> Self::Output {
    let rhs: u64 = rhs.try_into().expect("divisor must fit inside u64");

    match self.inner {
      DurationImpl::Zero => Duration::ZERO,
      DurationImpl::Nanos(nanos) => Duration::from_nanos(div_round(nanos, rhs)),
      DurationImpl::Ticks { ticks, frequency } => {
        Duration::from_ticks(div_round(ticks, rhs), frequency)
      },
    }
  }
}

// u64 division that rounds half-up while avoiding overflow and f64.
// Like regular division, panics on division by zero.
fn div_round(lhs: u64, rhs: u64) -> u64 {
  let div = lhs / rhs;
  let rem = lhs % rhs;

  div + u64::from(rem >= rhs - rem)
}

pub fn ticks_to_nanos(tsc_ticks: u64, tsc_freq_hz: u64) -> f64 {
  assert_ne!(tsc_freq_hz, 0);
  const NANOS_PER_SECOND: u128 = 1_000_000_000;

  // Using u128 to avoid any issues with intermediate values. Note that the TSC
  // counter can (in theory) start at _any_ u64 value.
  let ticks = u128::from(tsc_ticks);
  let ticks_per_second = u128::from(tsc_freq_hz);

  let numerator = ticks * NANOS_PER_SECOND;
  let whole_nanos = numerator / ticks_per_second;
  let nanos_remainder = numerator % ticks_per_second;
  let fractional_nanos = (nanos_remainder as f64) / (ticks_per_second as f64);

  // f64 limits us to 2^53 nanoseconds which is ~104 days. Should be more than
  // good enough for computing a function's runtime and even as a sum of
  // multiple runtime observations (what kind of benchmark runs > 100 days?).
  whole_nanos as f64 + fractional_nanos
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Instant {
  Ticks(u64),
  Std(std::time::Instant),
}

#[cfg(test)]
mod tests {
  use super::*;
  use speculoos::prelude::*;

  #[test]
  fn ticks_to_nanos_small_ticks() {
    // At 3.2 GHz, 1 tick == 0.3125 nanos
    assert_that(&ticks_to_nanos(124, 3_200_000_000)).is_close_to(38.75, 0.01);
    assert_that(&ticks_to_nanos(123, 3_200_000_000))
      .is_close_to(38.4375, 0.0001);
    assert_that(&ticks_to_nanos(122, 3_200_000_000)).is_close_to(38.125, 0.001);
    assert_that(&ticks_to_nanos(121, 3_200_000_000))
      .is_close_to(37.8125, 0.0001);
    assert_that(&ticks_to_nanos(120, 3_200_000_000)).is_close_to(37.5, 0.1);
    assert_that(&ticks_to_nanos(119, 3_200_000_000))
      .is_close_to(37.1875, 0.001);
  }

  #[test]
  fn ticks_to_nanos_big_ticks() {
    // At 3.2 GHz, 1 tick == 0.3125 nanos
    assert_eq!(
      ticks_to_nanos(267_597_764_222_976u64, 3_200_000_000),
      83_624_301_319_680f64
    );
    assert_that(&ticks_to_nanos(267_597_764_222_971, 3_200_000_000))
      .is_close_to(83_624_301_319_678.44, 0.01);
  }

  #[test]
  fn ticks_to_nanos_low_freq() {
    // Apple M2 uses a 24 MHz TSC frequency.
    // At 24 Mhz, 1 tick ~= 41.67 nanos
    assert_eq!(ticks_to_nanos(0, 24_000_000), 0.0);
    assert_that(&ticks_to_nanos(1, 24_000_000)).is_close_to(41.67, 0.01);
    assert_that(&ticks_to_nanos(2, 24_000_000)).is_close_to(83.34, 0.01);
  }

  #[test]
  fn from_nanos_conversion() {
    let duration = Duration::from_nanos(10);
    assert_eq!(duration.to_nanos(), 10.0);
    assert_eq!(duration.to_raw(), 10);
    assert_eq!(
      duration.to_std_duration(),
      std::time::Duration::from_nanos(10)
    );
  }

  #[test]
  fn from_ticks_conversion() {
    let duration = Duration::from_ticks(30, 3_000_000_000);
    assert_eq!(duration.to_nanos(), 10.0);
    assert_eq!(duration.to_raw(), 30);
    assert_eq!(
      duration.to_std_duration(),
      std::time::Duration::from_nanos(10)
    );
  }

  #[test]
  fn from_zero_nanos() {
    let duration = Duration::from_nanos(0);
    assert_eq!(duration.to_nanos(), 0.0);
    assert_eq!(duration.to_raw(), 0);
    assert_eq!(
      duration.to_std_duration(),
      std::time::Duration::from_nanos(0)
    );
    assert_eq!(duration, Duration::ZERO);
  }

  #[test]
  fn from_zero_ticks() {
    let duration = Duration::from_ticks(0, 3_000_000_000);
    assert_eq!(duration.to_nanos(), 0.0);
    assert_eq!(duration.to_raw(), 0);
    assert_eq!(
      duration.to_std_duration(),
      std::time::Duration::from_nanos(0)
    );
    assert_eq!(duration, Duration::ZERO);
  }

  #[test]
  fn addition_ticks() {
    let dur1 = Duration::from_ticks(10, 3_000_000_000);
    let dur2 = Duration::from_ticks(20, 3_000_000_000);
    assert_eq!((dur1 + dur2).to_raw(), 30);

    assert_eq!((dur1 + Duration::ZERO).to_raw(), 10);
    assert_eq!((Duration::ZERO + dur1).to_raw(), 10);
  }

  #[test]
  fn addition_nanos() {
    let dur1 = Duration::from_nanos(10);
    let dur2 = Duration::from_nanos(20);
    assert_eq!((dur1 + dur2).to_nanos(), 30.0);

    assert_eq!((dur1 + Duration::ZERO).to_nanos(), 10.0);
    assert_eq!((Duration::ZERO + dur1).to_nanos(), 10.0);
  }

  #[test]
  fn addition_mix() {
    let dur1 = Duration::from_nanos(10);
    let dur2 = Duration::from_ticks(20, 1_000_000_000);
    assert_eq!((dur1 + dur2).to_nanos(), 30.0);
    assert_eq!((dur2 + dur1).to_nanos(), 30.0);
  }

  #[test]
  fn subtraction_ticks() {
    let dur1 = Duration::from_ticks(10, 3_000_000_000);
    let dur2 = Duration::from_ticks(30, 3_000_000_000);
    assert_eq!((dur2 - dur1).to_raw(), 20);
    assert_eq!((dur1 - dur2).to_raw(), 0); // saturates!

    assert_eq!((dur1 - Duration::ZERO).to_raw(), 10);
    assert_eq!((Duration::ZERO - dur1).to_raw(), 0); // saturates!
  }

  #[test]
  fn subtraction_nanos() {
    let dur1 = Duration::from_nanos(10);
    let dur2 = Duration::from_nanos(30);
    assert_eq!((dur2 - dur1).to_nanos(), 20.0);
    assert_eq!((dur1 - dur2).to_nanos(), 0.0); // saturates!

    assert_eq!((dur1 - Duration::ZERO).to_nanos(), 10.0);
    assert_eq!((Duration::ZERO - dur1).to_nanos(), 0.0); // saturates!
  }

  #[test]
  fn subtraction_mix() {
    let dur1 = Duration::from_nanos(10);
    let dur2 = Duration::from_ticks(30, 1_000_000_000);
    assert_eq!((dur2 - dur1).to_nanos(), 20.0);

    let dur3 = Duration::from_nanos(50);
    assert_eq!((dur3 - dur2).to_nanos(), 20.0);
  }

  #[test]
  fn division_ticks() {
    let duration = Duration::from_ticks(7, 3_000_000_000);
    assert_eq!((duration / 3).to_raw(), 2);
    assert_eq!((duration / 4).to_raw(), 2);
    assert_eq!((duration / 5).to_raw(), 1);
    assert_eq!((duration / 8).to_raw(), 1);
    assert_eq!((duration / 15).to_raw(), 0);
  }

  #[test]
  fn division_nanos() {
    let duration = Duration::from_nanos(7);
    assert_eq!((duration / 3).to_raw(), 2);
    assert_eq!((duration / 4).to_raw(), 2);
    assert_eq!((duration / 5).to_raw(), 1);
    assert_eq!((duration / 8).to_raw(), 1);
    assert_eq!((duration / 15).to_raw(), 0);
  }

  #[test]
  fn ord_is_sensible() {
    assert!(Duration::ZERO < Duration::from_ticks(1, 100));
    assert!(Duration::ZERO < Duration::from_nanos(1));
    assert!(Duration::from_nanos(1) < Duration::from_nanos(2));
    assert!(Duration::from_ticks(1, 10) < Duration::from_ticks(2, 10));

    assert!(Duration::from_ticks(1, 1_000_000_000) < Duration::from_nanos(2));
    assert!(Duration::from_nanos(2) > Duration::from_ticks(1, 1_000_000_000));
  }

  #[test]
  fn eq_is_sensible() {
    assert!(Duration::ZERO == Duration::ZERO);
    assert!(Duration::from_nanos(1) == Duration::from_nanos(1));
    assert!(Duration::from_ticks(1, 10) == Duration::from_ticks(1, 10));

    assert!(Duration::from_ticks(1, 1_000_000_000) == Duration::from_nanos(1));
    assert!(Duration::from_ticks(0, 1_000_000_000) == Duration::from_nanos(0));
    assert!(Duration::from_ticks(0, 1_000_000_000) == Duration::ZERO);
    assert!(Duration::from_nanos(0) == Duration::ZERO);
  }

  #[test]
  fn storage_type_works() {
    assert_eq!(Duration::ZERO.storage_type(), StorageType::Zero);
    assert_eq!(Duration::from_nanos(1).storage_type(), StorageType::Nanos);
    assert_eq!(
      Duration::from_ticks(1, 10).storage_type(),
      StorageType::Ticks
    );
  }
}