span-timing 0.1.0

#![doc = include_str!("../README.md")]
#![cfg_attr(not(feature = "std"), no_std)]

#[cfg(all(
    not(feature = "std"),
    not(any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64"))
))]
compile_error!(
    "span-timing without the `std` feature supports only x86, x86_64, and aarch64 targets"
);

/// Receives measurements collected by [`timed_span!`].
///
/// The macro calls [`Self::increment_count`] when a timed scope begins and
/// [`Self::add_elapsed_ticks`] when it ends. Implementations may store those
/// values directly, aggregate them differently, or update additional metrics.
/// [`timing_entries!`] uses [`Self::INITIAL`] to create each element of a
/// generated static counter array.
///
/// # Example
///
/// This counter computes a running average on demand from its sample count and
/// total ticks:
///
/// ```
/// use std::sync::atomic::{AtomicU64, Ordering};
/// use span_timing::TimingCounter;
///
/// struct AverageCounter {
///     samples: AtomicU64,
///     total_ticks: AtomicU64,
/// }
///
/// impl AverageCounter {
///     fn average_ticks(&self) -> Option<u64> {
///         let samples = self.samples.load(Ordering::Relaxed);
///         (samples != 0).then(|| self.total_ticks.load(Ordering::Relaxed) / samples)
///     }
/// }
///
/// impl TimingCounter for AverageCounter {
///     const INITIAL: Self = Self {
///         samples: AtomicU64::new(0),
///         total_ticks: AtomicU64::new(0),
///     };
///
///     fn increment_count(&self) {
///         self.samples.fetch_add(1, Ordering::Relaxed);
///     }
///
///     fn add_elapsed_ticks(&self, ticks: u64) {
///         self.total_ticks.fetch_add(ticks, Ordering::Relaxed);
///     }
/// }
/// ```
pub trait TimingCounter {
    /// The const value used to initialize each element of a static counter collection.
    const INITIAL: Self;

    /// Records one invocation of the timed operation, before the timed scope runs.
    fn increment_count(&self);

    /// Records the elapsed processor-counter ticks or nanoseconds when the scope ends.
    fn add_elapsed_ticks(&self, elapsed_ticks: u64);
}

/// The standard atomic counter implementation for [`timed_span!`].
///
/// `count` records the number of timed spans, while `ticks` accumulates their
/// elapsed processor-counter ticks (or nanoseconds on unsupported architectures).
/// Use this type when a total, a count, and their derived average are sufficient.
#[cfg(target_has_atomic = "64")]
#[derive(Debug, Default)]
pub struct Counter {
    pub count: core::sync::atomic::AtomicU64,
    pub ticks: core::sync::atomic::AtomicU64,
}

#[cfg(target_has_atomic = "64")]
impl Counter {
    /// Creates a counter with both measurements set to zero.
    pub const fn new() -> Self {
        Self {
            count: core::sync::atomic::AtomicU64::new(0),
            ticks: core::sync::atomic::AtomicU64::new(0),
        }
    }

    /// Resets both measurements to zero.
    pub fn reset(&self) {
        self.count.store(0, core::sync::atomic::Ordering::Relaxed);
        self.ticks.store(0, core::sync::atomic::Ordering::Relaxed);
    }
}

#[cfg(target_has_atomic = "64")]
impl TimingCounter for Counter {
    const INITIAL: Self = Self::new();

    fn increment_count(&self) {
        self.count
            .fetch_add(1, core::sync::atomic::Ordering::Relaxed);
    }

    fn add_elapsed_ticks(&self, ticks: u64) {
        self.ticks
            .fetch_add(ticks, core::sync::atomic::Ordering::Relaxed);
    }
}

/// Declares an enum and optionally a static counter collection for timing entries.
///
/// The generated enum has `ALL`, `COUNT`, and `to_str` associated items. Adding
/// `static COUNTERS: [Counter];` after the enum generates a `[Counter; COUNT]` static,
/// initialized from [`TimingCounter::INITIAL`]. This keeps entry names, array length,
/// and reporting order in one declaration.
///
/// Without the `static` declaration, the macro only generates the enum and its helpers.
#[macro_export]
macro_rules! timing_entries {
    (
        $visibility:vis enum $name:ident {
            $($entry:ident $(= $value:expr)?),*
            $(,)?
        }
        $counter_visibility:vis static $counters:ident: [$counter_type:ty];
    ) => {
        $crate::timing_entries! {
            @entries
            $visibility enum $name {
                $($entry $(= $value)?),*
            }
        }

        $counter_visibility static $counters: [$counter_type; $name::COUNT] =
            [const { <$counter_type as $crate::TimingCounter>::INITIAL }; $name::COUNT];
    };
    (
        $visibility:vis enum $name:ident {
            $($entry:ident $(= $value:expr)?),*
            $(,)?
        }
    ) => {
        $crate::timing_entries! {
            @entries
            $visibility enum $name {
                $($entry $(= $value)?),*
            }
        }
    };
    (
        @entries
        $visibility:vis enum $name:ident {
            $($entry:ident $(= $value:expr)?),*
            $(,)?
        }
    ) => {
        #[derive(Clone, Copy)]
        $visibility enum $name {
            $($entry $(= $value)?),*
        }

        impl $name {
            pub const ALL: &[$name] = &[$($name::$entry),*];
            pub const COUNT: usize = $name::ALL.len();

            pub const fn to_str(&self) -> &'static str {
                match self {
                    $(
                        $name::$entry => stringify!($entry),
                    )*
                }
            }
        }
    };
}

/// Starts a timed span and returns its guard.
///
/// The first argument is an enum variant that can be cast to an index. The second is an
/// indexable collection whose entries implement [`TimingCounter`]. Bind the returned
/// guard for the scope to measure. When dropped, it adds the elapsed measurement, even
/// if the scope returns early or unwinds. On x86, x86_64, and aarch64 the elapsed value
/// is a processor-counter tick count; on other architectures it is elapsed nanoseconds.
#[macro_export]
macro_rules! timed_span {
    ($entry:expr, $counters:expr $(,)?) => {{
        let counter = &($counters)[($entry) as usize];
        $crate::TimingCounter::increment_count(counter);
        $crate::TimedSpanGuard::new(counter)
    }};
}

#[cfg(any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64"))]
mod clock {
    use super::TimingCounter;
    use core::arch::asm;

    #[cfg(target_arch = "aarch64")]
    #[inline]
    fn read_counter() -> u64 {
        let value: u64;
        unsafe {
            asm!("mrs {}, CNTVCT_EL0", out(reg) value, options(nostack, nomem));
        }
        value
    }

    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    #[inline]
    fn read_counter() -> u64 {
        let low: u32;
        let high: u32;
        unsafe {
            asm!(
                "rdtsc",
                out("eax") low,
                out("edx") high,
                options(nostack, nomem)
            );
        }
        ((high as u64) << 32) | low as u64
    }

    /// A scope guard that adds elapsed processor-counter ticks to an atomic counter.
    pub struct TimedSpanGuard<'a, C: TimingCounter + ?Sized> {
        start: u64,
        counter: &'a C,
    }

    impl<'a, C: TimingCounter + ?Sized> TimedSpanGuard<'a, C> {
        /// Starts timing and records elapsed ticks in `counter` when dropped.
        pub fn new(counter: &'a C) -> Self {
            Self {
                start: read_counter(),
                counter,
            }
        }
    }

    impl<C: TimingCounter + ?Sized> Drop for TimedSpanGuard<'_, C> {
        fn drop(&mut self) {
            self.counter.add_elapsed_ticks(read_counter() - self.start);
        }
    }
}

#[cfg(all(
    feature = "std",
    not(any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64"))
))]
mod clock {
    use super::TimingCounter;
    use std::time::Instant;

    /// A scope guard that adds elapsed nanoseconds to an atomic counter.
    pub struct TimedSpanGuard<'a, C: TimingCounter + ?Sized> {
        start: Instant,
        counter: &'a C,
    }

    impl<'a, C: TimingCounter + ?Sized> TimedSpanGuard<'a, C> {
        /// Starts timing and records elapsed nanoseconds in `counter` when dropped.
        pub fn new(counter: &'a C) -> Self {
            Self {
                start: Instant::now(),
                counter,
            }
        }
    }

    impl<C: TimingCounter + ?Sized> Drop for TimedSpanGuard<'_, C> {
        fn drop(&mut self) {
            self.counter
                .add_elapsed_ticks(self.start.elapsed().as_nanos() as u64);
        }
    }
}

pub use clock::TimedSpanGuard;

#[cfg(all(test, target_has_atomic = "64"))]
mod tests {
    use crate::{Counter as StandardCounter, TimingCounter};
    use core::sync::atomic::{AtomicU64, Ordering};

    timing_entries! {
        pub enum Entry {
            First,
            Second,
        }
        static COUNTERS: [Counter];
    }

    struct Counter {
        invocations: AtomicU64,
        elapsed: AtomicU64,
    }

    impl TimingCounter for Counter {
        const INITIAL: Self = Self {
            invocations: AtomicU64::new(0),
            elapsed: AtomicU64::new(0),
        };

        fn increment_count(&self) {
            self.invocations.fetch_add(1, Ordering::Relaxed);
        }

        fn add_elapsed_ticks(&self, elapsed_ticks: u64) {
            self.elapsed.fetch_add(elapsed_ticks, Ordering::Relaxed);
        }
    }

    #[test]
    fn standard_counter_records_and_resets_measurements() {
        let counter = StandardCounter::default();
        counter.increment_count();
        counter.add_elapsed_ticks(42);

        assert_eq!(counter.count.load(Ordering::Relaxed), 1);
        assert_eq!(counter.ticks.load(Ordering::Relaxed), 42);

        counter.reset();
        assert_eq!(counter.count.load(Ordering::Relaxed), 0);
        assert_eq!(counter.ticks.load(Ordering::Relaxed), 0);
    }

    #[test]
    fn declares_entries_and_records_a_span() {
        assert_eq!(Entry::ALL.len(), 2);
        assert_eq!(Entry::Second.to_str(), "Second");

        {
            let _timed_span_guard = timed_span!(Entry::First, COUNTERS);
            for value in 0..100_000 {
                core::hint::black_box(value);
            }
        }
        assert_eq!(
            COUNTERS[Entry::First as usize]
                .invocations
                .load(Ordering::Relaxed),
            1
        );
        assert_ne!(
            COUNTERS[Entry::First as usize]
                .elapsed
                .load(Ordering::Relaxed),
            0
        );
    }
}