1//! Types used to implement runtime argument support.
2
3use std::{
4 any::{Any, TypeId},
5 borrow::Cow,
6 mem, slice,
7 sync::OnceLock,
8};
9
10use crate::{util::ty::TypeCast, Bencher};
11
12/// Holds lazily-initialized runtime arguments to be passed into a benchmark.
13///
14/// `#[divan::bench]` stores this as a `__DIVAN_ARGS` global for each entry, and
15/// then at runtime it is initialized once by a closure that creates the usable
16/// `BenchArgsRunner`.
17pub struct BenchArgs {
18 args: OnceLock<ErasedArgsSlice>,
19}
20
21/// The result of making `BenchArgs` runnable from instantiating the arguments
22/// list and providing a typed benchmarking implementation.
23#[derive(Clone, Copy)]
24pub struct BenchArgsRunner {
25 args: &'static ErasedArgsSlice,
26 bench: fn(Bencher, &ErasedArgsSlice, arg_index: usize),
27}
28
29/// Type-erased `&'static [T]` that also stores names of the arguments.
30struct ErasedArgsSlice {
31 /// The start of `&[T]`.
32 args: *const (),
33
34 /// The start of `&[&'static str]`.
35 names: *const &'static str,
36
37 /// The number of arguments.
38 len: usize,
39
40 /// The ID of `T` to ensure correctness.
41 arg_type: TypeId,
42}
43
44// SAFETY: Raw pointers in `ErasedArgsSlice` are used in a thread-safe way, and
45// the argument type is required to be `Send + Sync` when initialized from the
46// iterator in `BenchArgs::runner`.
47unsafe impl Send for ErasedArgsSlice {}
48unsafe impl Sync for ErasedArgsSlice {}
49
50impl BenchArgs {
51 /// Creates an uninitialized instance.
52 pub const fn new() -> Self {
53 Self { args: OnceLock::new() }
54 }
55
56 /// Initializes `self` with the results of `make_args` and returns a
57 /// `BenchArgsRunner` that will execute the benchmarking closure.
58 pub fn runner<I, B>(
59 &'static self,
60 make_args: impl FnOnce() -> I,
61 arg_to_string: impl Fn(&I::Item) -> String,
62 _bench_impl: B,
63 ) -> BenchArgsRunner
64 where
65 I: IntoIterator,
66 I::Item: Any + Send + Sync,
67 B: FnOnce(Bencher, &I::Item) + Copy,
68 {
69 let args = self.args.get_or_init(|| {
70 let args_iter = make_args().into_iter();
71
72 // Reuse arguments for names if already a slice of strings.
73 //
74 // NOTE: We do this over `I::IntoIter` instead of `I` since it works
75 // for both slices and `slice::Iter`.
76 let args_strings: Option<&'static [&str]> =
77 args_iter.cast_ref::<slice::Iter<&str>>().map(|iter| iter.as_slice());
78
79 // Collect arguments into leaked slice.
80 //
81 // Leaking the collected `args` simplifies memory management, such
82 // as when reusing for `names`. We're leaking anyways since this is
83 // accessed via a global `OnceLock`.
84 //
85 // PERF: We could optimize this to reuse arguments when users
86 // provide slices. However, for slices its `Item` is a reference, so
87 // `slice::Iter<I::Item>` would never match here. To make this
88 // optimization, we would need to be able to get the referee type.
89 let args: &'static [I::Item] = Box::leak(args_iter.collect());
90
91 // Collect printable representations of arguments.
92 //
93 // PERF: We take multiple opportunities to reuse the provided
94 // arguments buffer or individual strings' buffers:
95 // - `&[&str]`
96 // - `IntoIterator<Item = &str>`
97 // - `IntoIterator<Item = String>`
98 // - `IntoIterator<Item = Box<str>>`
99 // - `IntoIterator<Item = Cow<str>>`
100 let names: &'static [&str] = 'names: {
101 // PERF: Reuse arguments strings slice.
102 if let Some(args) = args_strings {
103 break 'names args;
104 }
105
106 // PERF: Reuse our args slice allocation.
107 if let Some(args) = args.cast_ref::<&[&str]>() {
108 break 'names args;
109 }
110
111 Box::leak(
112 args.iter()
113 .map(|arg| -> &str {
114 // PERF: Reuse strings as-is.
115 if let Some(arg) = arg.cast_ref::<String>() {
116 return arg;
117 }
118 if let Some(arg) = arg.cast_ref::<Box<str>>() {
119 return arg;
120 }
121 if let Some(arg) = arg.cast_ref::<Cow<str>>() {
122 return arg;
123 }
124
125 // Default to `arg_to_string`, which will format via
126 // either `ToString` or `Debug`.
127 Box::leak(arg_to_string(arg).into_boxed_str())
128 })
129 .collect(),
130 )
131 };
132
133 ErasedArgsSlice {
134 // We `black_box` arguments to prevent the compiler from
135 // optimizing the benchmark for the provided values.
136 args: crate::black_box(args.as_ptr().cast()),
137 names: names.as_ptr(),
138 len: args.len(),
139 arg_type: TypeId::of::<I::Item>(),
140 }
141 });
142
143 BenchArgsRunner { args, bench: bench::<I::Item, B> }
144 }
145}
146
147impl Default for BenchArgs {
148 fn default() -> Self {
149 Self::new()
150 }
151}
152
153impl BenchArgsRunner {
154 #[inline]
155 pub(crate) fn bench(&self, bencher: Bencher, index: usize) {
156 (self.bench)(bencher, self.args, index)
157 }
158
159 #[inline]
160 pub(crate) fn arg_names(&self) -> &'static [&'static str] {
161 self.args.names()
162 }
163}
164
165impl ErasedArgsSlice {
166 /// Retrieves a slice of arguments if the type is `T`.
167 #[inline]
168 fn typed_args<T: Any>(&self) -> Option<&[T]> {
169 if self.arg_type == TypeId::of::<T>() {
170 // SAFETY: `BenchArgs::runner` guarantees storing `len` instances.
171 Some(unsafe { slice::from_raw_parts(self.args.cast(), self.len) })
172 } else {
173 None
174 }
175 }
176
177 /// Returns the arguments' names.
178 ///
179 /// Names are in the same order as args and thus their indices can be used
180 /// to reference arguments.
181 #[inline]
182 fn names(&self) -> &'static [&str] {
183 // SAFETY: `BenchArgs::runner` guarantees storing `len` names.
184 unsafe { slice::from_raw_parts(self.names, self.len) }
185 }
186}
187
188/// The `BenchArgsRunner.bench` implementation.
189fn bench<T, B>(bencher: Bencher, erased_args: &ErasedArgsSlice, arg_index: usize)
190where
191 T: Any,
192 B: FnOnce(Bencher, &T) + Copy,
193{
194 // We defer type checking until the benchmark is run to make safety of this
195 // function easier to audit. Checking here instead of in `BenchArgs::runner`
196 // is late but fine since this check will only fail due to a bug in Divan's
197 // macro code generation.
198
199 let Some(typed_args) = erased_args.typed_args::<T>() else {
200 type_mismatch::<T>();
201
202 // Reduce code size by using a separate function for each `T` instead of
203 // each benchmark closure.
204 #[cold]
205 #[inline(never)]
206 fn type_mismatch<T>() -> ! {
207 unreachable!("incorrect type '{}'", std::any::type_name::<T>())
208 }
209 };
210
211 // SAFETY: The closure is a ZST, so we can construct one out of thin air.
212 // This can be done multiple times without invoking a `Drop` destructor
213 // because it implements `Copy`.
214 let bench_impl: B = unsafe {
215 assert_eq!(size_of::<B>(), 0, "benchmark closure expected to be zero-sized");
216 mem::zeroed()
217 };
218
219 bench_impl(bencher, &typed_args[arg_index]);
220}
221
222#[cfg(test)]
223mod tests {
224 use super::*;
225
226 /// Test that optimizations for string items are applied.
227 mod optimizations {
228 use std::borrow::Borrow;
229
230 use super::*;
231
232 /// Tests that two slices contain the same exact strings.
233 fn test_eq_ptr<A: Borrow<str>, B: Borrow<str>>(a: &[A], b: &[B]) {
234 assert_eq!(a.len(), b.len());
235
236 for (a, b) in a.iter().zip(b) {
237 let a = a.borrow();
238 let b = b.borrow();
239 assert_eq!(a, b);
240 assert_eq!(a.as_ptr(), b.as_ptr());
241 }
242 }
243
244 /// Tests that `&[&str]` reuses the original slice for names.
245 #[test]
246 fn str_slice() {
247 static ARGS: BenchArgs = BenchArgs::new();
248 static ORIG_ARGS: &[&str] = &["a", "b"];
249
250 let runner = ARGS.runner(|| ORIG_ARGS, ToString::to_string, |_, _| {});
251
252 let typed_args: Vec<&str> =
253 runner.args.typed_args::<&&str>().unwrap().iter().copied().copied().collect();
254 let names = runner.arg_names();
255
256 // Test values.
257 assert_eq!(names, ORIG_ARGS);
258 assert_eq!(names, typed_args);
259
260 // Test addresses.
261 assert_eq!(names.as_ptr(), ORIG_ARGS.as_ptr());
262 assert_ne!(names.as_ptr(), typed_args.as_ptr());
263 }
264
265 /// Tests optimizing `IntoIterator<Item = &str>` to reuse the same
266 /// allocation for also storing argument names.
267 #[test]
268 fn str_array() {
269 static ARGS: BenchArgs = BenchArgs::new();
270
271 let runner = ARGS.runner(|| ["a", "b"], ToString::to_string, |_, _| {});
272
273 let typed_args = runner.args.typed_args::<&str>().unwrap();
274 let names = runner.arg_names();
275
276 // Test values.
277 assert_eq!(names, ["a", "b"]);
278 assert_eq!(names, typed_args);
279
280 // Test addresses.
281 assert_eq!(names.as_ptr(), typed_args.as_ptr());
282 }
283
284 /// Tests optimizing `IntoIterator<Item = String>` to reuse the same
285 /// allocation for also storing argument names.
286 #[test]
287 fn string_array() {
288 static ARGS: BenchArgs = BenchArgs::new();
289
290 let runner =
291 ARGS.runner(|| ["a".to_owned(), "b".to_owned()], ToString::to_string, |_, _| {});
292
293 let typed_args = runner.args.typed_args::<String>().unwrap();
294 let names = runner.arg_names();
295
296 assert_eq!(names, ["a", "b"]);
297 test_eq_ptr(names, typed_args);
298 }
299
300 /// Tests optimizing `IntoIterator<Item = Box<str>>` to reuse the same
301 /// allocation for also storing argument names.
302 #[test]
303 fn box_str_array() {
304 static ARGS: BenchArgs = BenchArgs::new();
305
306 let runner = ARGS.runner(
307 || ["a".to_owned().into_boxed_str(), "b".to_owned().into_boxed_str()],
308 ToString::to_string,
309 |_, _| {},
310 );
311
312 let typed_args = runner.args.typed_args::<Box<str>>().unwrap();
313 let names = runner.arg_names();
314
315 assert_eq!(names, ["a", "b"]);
316 test_eq_ptr(names, typed_args);
317 }
318
319 /// Tests optimizing `IntoIterator<Item = Cow<str>>` to reuse the same
320 /// allocation for also storing argument names.
321 #[test]
322 fn cow_str_array() {
323 static ARGS: BenchArgs = BenchArgs::new();
324
325 let runner = ARGS.runner(
326 || [Cow::Owned("a".to_owned()), Cow::Borrowed("b")],
327 ToString::to_string,
328 |_, _| {},
329 );
330
331 let typed_args = runner.args.typed_args::<Cow<str>>().unwrap();
332 let names = runner.arg_names();
333
334 assert_eq!(names, ["a", "b"]);
335 test_eq_ptr(names, typed_args);
336 }
337 }
338}
339