//! Handpicked stubs from [divan::entry](https://github.com/nvzqz/divan/blob/main/src/entry/mod.rs)
//! Necessary to be able to use the [divan::bench](https://docs.rs/divan/0.1.17/divan/attr.bench.html) macro without changing it too much

mod generic;

use std::{
    ptr,
    sync::{
        atomic::{AtomicPtr, Ordering as AtomicOrdering},
        LazyLock,
    },
};

use super::{
    bench::{BenchOptions, Bencher},
    BenchArgsRunner,
};

pub use generic::{EntryConst, EntryType, GenericBenchEntry};

/// Benchmark entries generated by `#[divan::bench]`.
///
/// Note: generic-type benchmark entries are instead stored in `GROUP_ENTRIES`
/// in `generic_benches`.
pub static BENCH_ENTRIES: EntryList<BenchEntry> = EntryList::root();

/// Group entries generated by `#[divan::bench_group]`.
pub static GROUP_ENTRIES: EntryList<GroupEntry> = EntryList::root();

/// Determines how the benchmark entry is run.
#[derive(Clone, Copy)]
pub enum BenchEntryRunner {
    /// Benchmark without arguments.
    Plain(fn(Bencher)),
    /// Benchmark with runtime arguments.
    Args(fn() -> BenchArgsRunner),
}

/// Compile-time entry for a benchmark, generated by `#[divan::bench]`.
pub struct BenchEntry {
    /// Entry metadata.
    pub meta: EntryMeta,

    /// The benchmarking function.
    pub bench: BenchEntryRunner,
}

/// Compile-time entry for a benchmark group, generated by
/// `#[divan::bench_group]` or a generic-type `#[divan::bench]`.
pub struct GroupEntry {
    /// Entry metadata.
    pub meta: EntryMeta,

    /// Generic `#[divan::bench]` entries.
    ///
    /// This is two-dimensional to make code generation simpler. The outer
    /// dimension corresponds to types and the inner dimension corresponds to
    /// constants.
    pub generic_benches: Option<&'static [&'static [GenericBenchEntry]]>,
}

impl GroupEntry {
    pub(crate) fn generic_benches_iter(&self) -> impl Iterator<Item = &'static GenericBenchEntry> {
        self.generic_benches
            .unwrap_or_default()
            .iter()
            .flat_map(|benches| benches.iter())
    }
}

/// `BenchEntry` or `GenericBenchEntry`.
#[derive(Clone, Copy)]
pub(crate) enum AnyBenchEntry<'a> {
    Bench(&'a BenchEntry),
    GenericBench(&'a GenericBenchEntry),
}

/// Metadata common to `#[divan::bench]` and `#[divan::bench_group]`.
pub struct EntryMeta {
    /// The entry's display name.
    pub display_name: &'static str,

    /// The entry's original name.
    ///
    /// This is used to find a `GroupEntry` for a `BenchEntry`.
    pub raw_name: &'static str,

    /// The entry's raw `module_path!()`.
    pub module_path: &'static str,

    /// Where the entry was defined.
    pub location: EntryLocation,
    /// Configures the benchmarker via attribute options.
    pub bench_options: Option<LazyLock<BenchOptions>>,
}

impl EntryMeta {
    #[inline]
    pub(crate) fn module_path_components<'a>(&self) -> impl Iterator<Item = &'a str> {
        self.module_path.split("::")
    }
}

/// Where an entry is located.
#[derive(Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord)]
#[allow(missing_docs)]
pub struct EntryLocation {
    pub file: &'static str,
    pub line: u32,
    pub col: u32,
}

/// Linked list of entries.
///
/// This is implemented in a thread-safe way despite the fact that constructors
/// are run single-threaded.
pub struct EntryList<T: 'static> {
    entry: Option<&'static T>,
    next: AtomicPtr<Self>,
}

impl<T> EntryList<T> {
    pub(crate) const fn root() -> Self {
        Self {
            entry: None,
            next: AtomicPtr::new(ptr::null_mut()),
        }
    }

    /// Dereferences the `next` pointer.
    #[inline]
    fn next(&self) -> Option<&Self> {
        // SAFETY: `next` is only assigned by `push`, which always receives a
        // 'static lifetime.
        unsafe { self.next.load(AtomicOrdering::Relaxed).as_ref() }
    }
}

// Externally used by macros or tests.
#[allow(missing_docs)]
impl<T> EntryList<T> {
    #[inline]
    pub const fn new(entry: &'static T) -> Self {
        Self {
            entry: Some(entry),
            next: AtomicPtr::new(ptr::null_mut()),
        }
    }

    /// Creates an iterator over entries in `self`.
    #[inline]
    pub fn iter(&self) -> impl Iterator<Item = &T> {
        let mut list = Some(self);
        std::iter::from_fn(move || -> Option<Option<&T>> {
            let current = list?;
            list = current.next();
            Some(current.entry.as_ref().copied())
        })
        .flatten()
    }

    /// Inserts `other` to the front of the list.
    ///
    /// # Safety
    ///
    /// This function must be safe to call before `main`.
    #[inline]
    pub fn push(&'static self, other: &'static Self) {
        let mut old_next = self.next.load(AtomicOrdering::Relaxed);
        loop {
            // Each publicly-created instance has `list.next` be null, so we can
            // simply store `self.next` there.
            other.next.store(old_next, AtomicOrdering::Release);

            // SAFETY: The content of `other` can already be seen, so we don't
            // need to strongly order reads into it.
            let other = other as *const Self as *mut Self;
            match self.next.compare_exchange_weak(
                old_next,
                other,
                AtomicOrdering::AcqRel,
                AtomicOrdering::Acquire,
            ) {
                // Successfully wrote our thread's value to the list.
                Ok(_) => return,

                // Lost the race, store winner's value in `other.next`.
                Err(new) => old_next = new,
            }
        }
    }
}

impl<'a> AnyBenchEntry<'a> {
    /// Returns this entry's benchmark runner.
    #[inline]
    pub fn bench_runner(self) -> &'a BenchEntryRunner {
        match self {
            Self::Bench(BenchEntry { bench, .. })
            | Self::GenericBench(GenericBenchEntry { bench, .. }) => bench,
        }
    }

    #[inline]
    pub fn meta(self) -> &'a EntryMeta {
        match self {
            Self::Bench(entry) => &entry.meta,
            Self::GenericBench(entry) => &entry.group.meta,
        }
    }

    #[inline]
    pub fn display_name(self) -> &'a str {
        match self {
            Self::Bench(entry) => entry.meta.display_name,
            Self::GenericBench(entry) => entry.display_name(),
        }
    }
}
