feat: manage type generics in codspeed instrumentation

Author: GuillaumeLagrange

crates/divan_compat/examples/benches/math.rs

@@ -91,7 +91,6 @@ mod fibonacci {
 
     // Will be ignored in instrumented mode as we do not support type generics yet
     // O(n)
-    #[cfg(not(codspeed))]
     #[divan::bench(
         types = [BTreeMap<u64, u64>, HashMap<u64, u64>],
         args = VALUES,

crates/divan_compat/src/compat/entry/generic.rs

@@ -0,0 +1,137 @@
+use std::{
+    any::{Any, TypeId},
+    cmp::Ordering,
+    mem::ManuallyDrop,
+    sync::OnceLock,
+};
+
+use super::{BenchEntryRunner, GroupEntry};
+
+// use crate::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 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!(),
+        }
+    }
+}
+
+/// 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`.
+#[allow(dead_code)]
+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 {}
+
+#[allow(dead_code)]
+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()),
+        }
+    }
+
+    /// [`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())
+        })
+    }
+}

crates/divan_compat/src/compat/entry.rs renamed to crates/divan_compat/src/compat/entry/mod.rs

@@ -1,5 +1,8 @@
 //! 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::{
@@ -13,12 +16,17 @@ use super::{
     BenchArgsRunner,
 };
 
+pub use generic::{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 {
@@ -37,6 +45,36 @@ pub struct BenchEntry {
     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.
@@ -152,3 +190,30 @@ impl<T> EntryList<T> {
         }
     }
 }
+
+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(),
+        }
+    }
+}

crates/divan_compat/src/compat/mod.rs

@@ -3,7 +3,10 @@
 pub mod __private {
     pub use super::{
         bench::{BenchArgs, BenchOptions},
-        entry::{BenchEntry, BenchEntryRunner, EntryList, EntryLocation, EntryMeta, BENCH_ENTRIES},
+        entry::{
+            BenchEntry, BenchEntryRunner, EntryList, EntryLocation, EntryMeta, EntryType,
+            GenericBenchEntry, GroupEntry, BENCH_ENTRIES, GROUP_ENTRIES,
+        },
     };
 
     pub use divan::__private::{Arg, ToStringHelper};
@@ -18,12 +21,25 @@ use std::{cell::RefCell, rc::Rc};
 
 pub use bench::*;
 use codspeed::codspeed::CodSpeed;
+use entry::AnyBenchEntry;
 
 pub fn main() {
     // Outlined steps of original divan::main and their equivalent in codspeed instrumented mode
     // 1. Get registered entries
-    // TODO: Manage bench groups
-    let bench_entries = &entry::BENCH_ENTRIES;
+    let group_entries = &entry::GROUP_ENTRIES;
+
+    let generic_bench_entries = group_entries.iter().flat_map(|group| {
+        group
+            .generic_benches_iter()
+            .map(AnyBenchEntry::GenericBench)
+    });
+
+    let bench_entries = entry::BENCH_ENTRIES
+        .iter()
+        .map(AnyBenchEntry::Bench)
+        .chain(generic_bench_entries);
+
+    // TODO: Manage non generic bench groups
 
     // 2. Build an execution tree
     // No need, we do not manage detailed tree printing like original divan, and we extract
@@ -35,25 +51,30 @@ pub fn main() {
 
     // 4. Scan the tree and execute benchmarks
     let codspeed = Rc::new(RefCell::new(CodSpeed::new()));
-    for entry in bench_entries.iter() {
-        let entry_uri = uri::generate(entry.meta.display_name, &entry.meta);
+    for entry in bench_entries {
+        let runner = entry.bench_runner();
+        let meta = entry.meta();
 
-        if let Some(options) = &entry.meta.bench_options.as_ref() {
+        if let Some(options) = &meta.bench_options {
             if let Some(true) = options.ignore {
-                println!("Skipped: {}", entry_uri);
+                let uri = uri::generate(&entry, entry.display_name());
+                println!("Skipped: {}", uri);
                 continue;
             }
         }
-        match entry.bench {
+        match runner {
             entry::BenchEntryRunner::Plain(bench_fn) => {
-                bench_fn(bench::Bencher::new(&codspeed, entry_uri));
+                let uri = uri::generate(&entry, entry.display_name());
+
+                bench_fn(bench::Bencher::new(&codspeed, uri));
             }
             entry::BenchEntryRunner::Args(bench_runner) => {
                 let bench_runner = bench_runner();
 
                 for (arg_index, arg_name) in bench_runner.arg_names().iter().enumerate() {
-                    let entry_name_with_arg = uri::append_arg(&entry_uri, arg_name);
-                    let bencher = bench::Bencher::new(&codspeed, entry_name_with_arg);
+                    let uri = uri::generate(&entry, arg_name);
+
+                    let bencher = bench::Bencher::new(&codspeed, uri);
 
                     bench_runner.bench(bencher, arg_index);
                 }