span-timing 0.1.1

#![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] = &{
                let declared: [$name; 0 $(+ { let _ = $name::$entry; 1 })*] =
                    [$($name::$entry),*];
                let mut ordered = declared;
                let mut index = 0;

                while index < declared.len() {
                    let entry = declared[index];
                    ordered[entry as usize] = entry;
                    index += 1;
                }

                ordered
            };
            pub const COUNT: usize = $name::ALL.len();

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

        $(
            const _: $name = $name::ALL[$name::$entry as usize];
        )*
    };
}

/// 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
        );
    }

    #[test]
    fn declares_entries_with_in_bounds_explicit_discriminants() {
        timing_entries! {
            enum ExplicitEntry {
                Second = 1,
                First = 0,
            }
        }

        assert_eq!(ExplicitEntry::COUNT, 2);
        assert_eq!(ExplicitEntry::First as usize, 0);
        assert_eq!(ExplicitEntry::Second as usize, 1);
        assert_eq!(
            ExplicitEntry::ALL[ExplicitEntry::First as usize].to_str(),
            "First"
        );
        assert_eq!(
            ExplicitEntry::ALL[ExplicitEntry::Second as usize].to_str(),
            "Second"
        );
    }
}

/// These tests use cargo directly to validate a compile failure.  Other alternatives are
/// 1. a doc test, but I didn't want to pollute the docs or,
/// 2. the `trybuild` crate, but I didn't want to rely on the exact compiler error message.
#[cfg(all(test, feature = "std"))]
mod compile_fail_tests {
    use std::{
        env, fs,
        process::{self, Command},
        time::{SystemTime, UNIX_EPOCH},
    };

    #[test]
    fn rejects_entries_with_out_of_bounds_discriminants() {
        let unique = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("system clock is before the Unix epoch")
            .as_nanos();
        let test_dir = env::temp_dir().join(format!("span-timing-oob-{unique}-{}", process::id()));
        let manifest_dir =
            env::var("CARGO_MANIFEST_DIR").expect("Cargo did not set CARGO_MANIFEST_DIR");

        fs::create_dir_all(test_dir.join("src")).expect("failed to create temporary test crate");
        fs::write(
            test_dir.join("Cargo.toml"),
            format!(
                "[package]\nname = \"timing-entries-oob\"\nversion = \"0.0.0\"\nedition = \"2024\"\n\n[dependencies]\nspan-timing = {{ path = {manifest_dir:?} }}\n"
            ),
        )
        .expect("failed to write temporary manifest");
        fs::write(
            test_dir.join("src/main.rs"),
            "use span_timing::timing_entries;\n\ntiming_entries! {\n    enum Oob {\n        Foo = 42,\n    }\n}\n\nfn main() {}\n",
        )
        .expect("failed to write temporary source");

        let output = Command::new(env::var("CARGO").unwrap_or_else(|_| "cargo".into()))
            .args(["check", "--offline"])
            .current_dir(&test_dir)
            .output()
            .expect("failed to run cargo check");
        let _ = fs::remove_dir_all(&test_dir);

        assert!(
            !output.status.success(),
            "out-of-bounds discriminant unexpectedly compiled:\n{}",
            String::from_utf8_lossy(&output.stdout)
        );
    }
}