codspeed-divan-compat 3.0.1

use crate::compat_divan::compat::util::ty::TypeCast;

use core::slice;
use std::{
    any::{Any, TypeId},
    borrow::Cow,
    mem,
    sync::OnceLock,
};

use super::Bencher;

/// Holds lazily-initialized runtime arguments to be passed into a benchmark.
///
/// `#[divan::bench]` stores this as a `__DIVAN_ARGS` global for each entry, and
/// then at runtime it is initialized once by a closure that creates the usable
/// `BenchArgsRunner`.
pub struct BenchArgs {
    args: OnceLock<ErasedArgsSlice>,
}

/// Type-erased `&'static [T]` that also stores names of the arguments.
struct ErasedArgsSlice {
    /// The start of `&[T]`.
    args: *const (),

    /// The start of `&[&'static str]`.
    names: *const &'static str,

    /// The number of arguments.
    len: usize,

    /// The ID of `T` to ensure correctness.
    arg_type: TypeId,
}

// SAFETY: Raw pointers in `ErasedArgsSlice` are used in a thread-safe way, and
// the argument type is required to be `Send + Sync` when initialized from the
// iterator in `BenchArgs::runner`.
unsafe impl Send for ErasedArgsSlice {}
unsafe impl Sync for ErasedArgsSlice {}

/// The result of making `BenchArgs` runnable from instantiating the arguments
/// list and providing a typed benchmarking implementation.
#[derive(Clone, Copy)]
pub struct BenchArgsRunner {
    args: &'static ErasedArgsSlice,
    bench: fn(Bencher, &ErasedArgsSlice, arg_index: usize),
}

impl BenchArgs {
    pub const fn new() -> Self {
        Self {
            args: OnceLock::new(),
        }
    }

    pub fn runner<I, B>(
        &'static self,
        make_args: impl FnOnce() -> I,
        arg_to_string: impl Fn(&I::Item) -> String,
        _bench_impl: B,
    ) -> BenchArgsRunner
    where
        I: IntoIterator,
        I::Item: Any + Send + Sync,
        B: FnOnce(Bencher, &I::Item) + Copy,
    {
        let args = self.args.get_or_init(|| {
            let args_iter = make_args().into_iter();

            // Reuse arguments for names if already a slice of strings.
            //
            // NOTE: We do this over `I::IntoIter` instead of `I` since it works
            // for both slices and `slice::Iter`.
            let args_strings: Option<&'static [&str]> = args_iter
                .cast_ref::<slice::Iter<&str>>()
                .map(|iter| iter.as_slice());

            // Collect arguments into leaked slice.
            //
            // Leaking the collected `args` simplifies memory management, such
            // as when reusing for `names`. We're leaking anyways since this is
            // accessed via a global `OnceLock`.
            //
            // PERF: We could optimize this to reuse arguments when users
            // provide slices. However, for slices its `Item` is a reference, so
            // `slice::Iter<I::Item>` would never match here. To make this
            // optimization, we would need to be able to get the referee type.
            let args: &'static [I::Item] = Box::leak(args_iter.collect());

            // Collect printable representations of arguments.
            //
            // PERF: We take multiple opportunities to reuse the provided
            // arguments buffer or individual strings' buffers:
            // - `&[&str]`
            // - `IntoIterator<Item = &str>`
            // - `IntoIterator<Item = String>`
            // - `IntoIterator<Item = Box<str>>`
            // - `IntoIterator<Item = Cow<str>>`
            let names: &'static [&str] = 'names: {
                // PERF: Reuse arguments strings slice.
                if let Some(args) = args_strings {
                    break 'names args;
                }

                // PERF: Reuse our args slice allocation.
                if let Some(args) = args.cast_ref::<&[&str]>() {
                    break 'names args;
                }

                Box::leak(
                    args.iter()
                        .map(|arg| -> &str {
                            // PERF: Reuse strings as-is.
                            if let Some(arg) = arg.cast_ref::<String>() {
                                return arg;
                            }
                            if let Some(arg) = arg.cast_ref::<Box<str>>() {
                                return arg;
                            }
                            if let Some(arg) = arg.cast_ref::<Cow<str>>() {
                                return arg;
                            }

                            // Default to `arg_to_string`, which will format via
                            // either `ToString` or `Debug`.
                            Box::leak(arg_to_string(arg).into_boxed_str())
                        })
                        .collect(),
                )
            };

            ErasedArgsSlice {
                // We `black_box` arguments to prevent the compiler from
                // optimizing the benchmark for the provided values.
                args: crate::black_box(args.as_ptr().cast()),
                names: names.as_ptr(),
                len: args.len(),
                arg_type: TypeId::of::<I::Item>(),
            }
        });

        BenchArgsRunner {
            args,
            bench: bench::<I::Item, B>,
        }
    }
}

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

impl BenchArgsRunner {
    #[inline]
    pub(crate) fn bench(&self, bencher: Bencher, index: usize) {
        (self.bench)(bencher, self.args, index)
    }

    #[inline]
    pub(crate) fn arg_names(&self) -> &'static [&'static str] {
        self.args.names()
    }
}

impl ErasedArgsSlice {
    /// Retrieves a slice of arguments if the type is `T`.
    #[inline]
    fn typed_args<T: Any>(&self) -> Option<&[T]> {
        if self.arg_type == TypeId::of::<T>() {
            // SAFETY: `BenchArgs::runner` guarantees storing `len` instances.
            Some(unsafe { slice::from_raw_parts(self.args.cast(), self.len) })
        } else {
            None
        }
    }

    /// Returns the arguments' names.
    ///
    /// Names are in the same order as args and thus their indices can be used
    /// to reference arguments.
    #[inline]
    fn names(&self) -> &'static [&str] {
        // SAFETY: `BenchArgs::runner` guarantees storing `len` names.
        unsafe { slice::from_raw_parts(self.names, self.len) }
    }
}

/// The `BenchArgsRunner.bench` implementation.
fn bench<T, B>(bencher: Bencher, erased_args: &ErasedArgsSlice, arg_index: usize)
where
    T: Any,
    B: FnOnce(Bencher, &T) + Copy,
{
    // We defer type checking until the benchmark is run to make safety of this
    // function easier to audit. Checking here instead of in `BenchArgs::runner`
    // is late but fine since this check will only fail due to a bug in Divan's
    // macro code generation.

    let Some(typed_args) = erased_args.typed_args::<T>() else {
        type_mismatch::<T>();

        // Reduce code size by using a separate function for each `T` instead of
        // each benchmark closure.
        #[cold]
        #[inline(never)]
        fn type_mismatch<T>() -> ! {
            unreachable!("incorrect type '{}'", std::any::type_name::<T>())
        }
    };

    // SAFETY: The closure is a ZST, so we can construct one out of thin air.
    // This can be done multiple times without invoking a `Drop` destructor
    // because it implements `Copy`.
    let bench_impl: B = unsafe {
        assert_eq!(
            size_of::<B>(),
            0,
            "benchmark closure expected to be zero-sized"
        );
        mem::zeroed()
    };

    bench_impl(bencher, &typed_args[arg_index]);
}