generic.rs

use std::{
    any::{Any, TypeId},
    cmp::Ordering,
    mem::ManuallyDrop,
    sync::OnceLock,
};

use crate::{
    entry::{BenchEntryRunner, GroupEntry},
    util::sort::natural_cmp,
};

/// Compile-time entry for a generic benchmark function, generated by
/// `#[divan::bench]`.
///
/// Unlike `BenchEntry`, this is for a specific generic type or `const`.
///
/// Although this type contains trivially-`Copy` data, it *should not* implement
/// `Clone` because the memory address of each instance is used to determine the
/// relative order in `GroupEntry.generic_benches` when sorting benchmarks by
/// location.
pub struct GenericBenchEntry {
    /// The associated group, for entry metadata.
    pub group: &'static GroupEntry,

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

    /// A generic type.
    pub ty: Option<EntryType>,

    /// A `const` value and associated data.
    pub const_value: Option<EntryConst>,
}

impl GenericBenchEntry {
    pub(crate) fn raw_name(&self) -> &str {
        match (&self.ty, &self.const_value) {
            (_, Some(const_value)) => const_value.name(),
            (Some(ty), None) => ty.raw_name(),
            (None, None) => unreachable!(),
        }
    }

    pub(crate) fn display_name(&self) -> &str {
        match (&self.ty, &self.const_value) {
            (_, Some(const_value)) => const_value.name(),
            (Some(ty), None) => ty.display_name(),
            (None, None) => unreachable!(),
        }
    }

    pub(crate) fn path_components(&self) -> impl Iterator<Item = &str> {
        let module_path = self.group.meta.module_path_components();

        // Generic benchmarks consider their group's raw name to be the path
        // component after the module path.
        let group_component = self.group.meta.raw_name;

        // If this is a generic const benchmark with generic types, the generic
        // types are considered to be the parent of the const values.
        let type_component = if self.const_value.is_some() {
            // FIXME: Switch back to `raw_name` once we have a way to insert
            // this `display_name` into `EntryTree::Parent`. The current
            // approach allows different types with the same name to become the
            // same `EntryTree::Parent`.
            self.ty.as_ref().map(|ty| ty.display_name())
        } else {
            None
        };

        module_path.chain(Some(group_component)).chain(type_component)
    }
}

/// Generic type instantiation.
pub struct EntryType {
    /// [`std::any::type_name`].
    get_type_name: fn() -> &'static str,

    /// [`std::any::TypeId::of`].
    #[allow(dead_code)]
    get_type_id: fn() -> TypeId,
}

impl EntryType {
    /// Creates an instance for the given type.
    pub const fn new<T: Any>() -> Self {
        Self { get_type_name: std::any::type_name::<T>, get_type_id: TypeId::of::<T> }
    }

    pub(crate) fn raw_name(&self) -> &'static str {
        (self.get_type_name)()
    }

    pub(crate) fn display_name(&self) -> &'static str {
        let mut type_name = self.raw_name();

        // Remove module components in type name.
        while let Some((prev, next)) = type_name.split_once("::") {
            // Do not go past generic type boundary.
            if prev.contains('<') {
                break;
            }
            type_name = next;
        }

        type_name
    }
}

/// A reference to a `const` as a `&'static T`.
pub struct EntryConst {
    /// `&'static T`.
    value: *const (),

    /// [`PartialOrd::partial_cmp`].
    partial_cmp: unsafe fn(*const (), *const ()) -> Option<Ordering>,

    /// [`ToString::to_string`].
    to_string: unsafe fn(*const ()) -> String,

    /// Cached `to_string` result.
    cached_string: ManuallyDrop<OnceLock<&'static str>>,
}

// SAFETY: `T: Send + Sync`.
unsafe impl Send for EntryConst {}
unsafe impl Sync for EntryConst {}

impl EntryConst {
    /// Creates entry data for a `const` values.
    pub const fn new<T>(value: &'static T) -> Self
    where
        T: PartialOrd + ToString + Send + Sync,
    {
        unsafe fn partial_cmp<T: PartialOrd>(a: *const (), b: *const ()) -> Option<Ordering> {
            T::partial_cmp(&*a.cast(), &*b.cast())
        }

        unsafe fn to_string<T: ToString>(value: *const ()) -> String {
            T::to_string(&*value.cast())
        }

        Self {
            value: value as *const T as *const (),
            partial_cmp: partial_cmp::<T>,
            to_string: to_string::<T>,
            cached_string: ManuallyDrop::new(OnceLock::new()),
        }
    }

    /// Returns [`PartialOrd::partial_cmp`] ordering if `<` or `>, falling back
    /// to comparing [`ToString::to_string`] otherwise.
    pub(crate) fn cmp_name(&self, other: &Self) -> Ordering {
        if self.partial_cmp as usize == other.partial_cmp as usize {
            // SAFETY: Both constants have the same comparison function, so they
            // must be the same type.
            if let Some(ordering) = unsafe { (self.partial_cmp)(self.value, other.value) } {
                if !ordering.is_eq() {
                    return ordering;
                }
            }
        }

        // Fallback to name comparison.
        natural_cmp(self.name(), other.name())
    }

    /// [`ToString::to_string`].
    #[inline]
    pub(crate) fn name(&self) -> &str {
        self.cached_string.get_or_init(|| {
            // SAFETY: The function is guaranteed to call `T::to_string`.
            let string = unsafe { (self.to_string)(self.value) };

            Box::leak(string.into_boxed_str())
        })
    }
}