1use std::time::{Duration, Instant};
2
3pub use core::hint::black_box as slain;
4
5pub struct Bencher {
6 iters: u64,
7 ticks: u64,
8}
9
10impl Bencher {
11 pub fn iter<R>(&mut self, mut f: impl FnMut() -> R) {
12 let t0 = clock::now();
13 for _ in 0..self.iters {
14 slain(f());
15 }
16 self.ticks = clock::now().wrapping_sub(t0);
17 }
18}
19
20type Arm<'a> = (String, Box<dyn FnMut(&mut Bencher) + 'a>);
21
22pub struct Group<'a> {
23 name: String,
24 arms: Vec<Arm<'a>>,
25 cfg: Config,
26}
27
28#[derive(Clone, Copy)]
29struct Config {
30 warm: Duration,
31 sample: Duration,
32 rounds: usize,
33}
34
35impl Default for Config {
36 fn default() -> Self {
37 Self { warm: Duration::from_millis(300), sample: Duration::from_micros(2000), rounds: 41 }
38 }
39}
40
41impl<'a> Group<'a> {
42 pub fn bench_function(&mut self, name: impl Into<String>, f: impl FnMut(&mut Bencher) + 'a) -> &mut Self {
43 self.arms.push((name.into(), Box::new(f)));
44 self
45 }
46
47 pub fn finish(mut self) {
48 if self.arms.is_empty() {
49 return;
50 }
51 let report = run(&self.name, &mut self.arms, self.cfg);
52 print!("{report}");
53 }
54}
55
56#[derive(Default)]
57pub struct Harness {
58 filter: Option<String>,
59}
60
61impl Harness {
62 #[must_use]
63 pub fn from_args() -> Self {
64 pin_and_raise();
65 Self { filter: std::env::args().skip(1).find(|a| !a.starts_with('-')) }
66 }
67
68 pub fn benchmark_group<'a>(&mut self, name: impl Into<String>) -> Group<'a> {
69 let name = name.into();
70 let skip = self.filter.as_ref().is_some_and(|f| !name.contains(f.as_str()));
71 Group { name, arms: Vec::new(), cfg: if skip { Config { rounds: 0, ..Config::default() } } else { Config::default() } }
72 }
73}
74
75#[macro_export]
76macro_rules! bench_main {
77 ($($f:path),+ $(,)?) => {
78 fn main() {
79 let mut h = $crate::bench::Harness::from_args();
80 $($f(&mut h);)+
81 }
82 };
83}
84
85struct Measured {
86 name: String,
87 ns: Vec<f64>,
88}
89
90fn run(group: &str, arms: &mut [Arm<'_>], cfg: Config) -> String {
91 if cfg.rounds == 0 {
92 return String::new();
93 }
94 let mut slots: Vec<usize> = (0..arms.len()).collect();
95 slots.push(0);
96 let iters: Vec<u64> = (0..arms.len()).map(|a| calibrate(&mut arms[a].1, cfg.sample)).collect();
97 let t = Instant::now();
98 while t.elapsed() < cfg.warm {
99 for (a, arm) in arms.iter_mut().enumerate() {
100 sample(&mut arm.1, iters[a]);
101 }
102 }
103 let overhead = timer_overhead();
104 let mut data: Vec<Measured> = slots
105 .iter()
106 .enumerate()
107 .map(|(s, &a)| Measured { name: if s >= arms.len() { format!("{} (A/A control)", arms[a].0) } else { arms[a].0.clone() }, ns: Vec::with_capacity(cfg.rounds) })
108 .collect();
109 for round in 0..cfg.rounds {
110 for k in 0..slots.len() {
111 let s = (k + round) % slots.len();
112 let a = slots[s];
113 let ticks = sample(&mut arms[a].1, iters[a]).saturating_sub(overhead);
114 data[s].ns.push(clock::to_ns(ticks) / iters[a] as f64);
115 }
116 }
117 report(group, &data)
118}
119
120fn calibrate(f: &mut Box<dyn FnMut(&mut Bencher) + '_>, target: Duration) -> u64 {
121 let mut iters = 1u64;
122 loop {
123 let t = Instant::now();
124 sample(f, iters);
125 let e = t.elapsed();
126 if e >= target || iters >= 1 << 40 {
127 return iters;
128 }
129 let scale = (target.as_secs_f64() / e.as_secs_f64().max(1e-9)).clamp(1.5, 100.0);
130 iters = ((iters as f64) * scale).ceil() as u64;
131 }
132}
133
134fn sample(f: &mut Box<dyn FnMut(&mut Bencher) + '_>, iters: u64) -> u64 {
135 let mut b = Bencher { iters, ticks: 0 };
136 f(&mut b);
137 b.ticks
138}
139
140fn timer_overhead() -> u64 {
141 let mut best = u64::MAX;
142 for _ in 0..1000 {
143 let t0 = clock::now();
144 best = best.min(clock::now().wrapping_sub(t0));
145 }
146 best
147}
148
149fn median(v: &mut [f64]) -> f64 {
150 v.sort_by(f64::total_cmp);
151 let n = v.len();
152 if n == 0 {
153 f64::NAN
154 } else if n % 2 == 1 {
155 v[n / 2]
156 } else {
157 (v[n / 2 - 1] + v[n / 2]) / 2.0
158 }
159}
160
161fn bootstrap(x: &[f64]) -> (f64, f64) {
162 let mut state = 0x9E37_79B9_7F4A_7C15u64;
163 let mut next = || {
164 state ^= state << 13;
165 state ^= state >> 7;
166 state ^= state << 17;
167 state
168 };
169 let n = x.len();
170 let mut meds: Vec<f64> = (0..2000)
171 .map(|_| {
172 let mut s: Vec<f64> = (0..n).map(|_| x[(next() % n as u64) as usize]).collect();
173 median(&mut s)
174 })
175 .collect();
176 meds.sort_by(f64::total_cmp);
177 (meds[50], meds[1949])
178}
179
180fn report(group: &str, data: &[Measured]) -> String {
181 use std::fmt::Write;
182 let mut out = String::new();
183 let base = &data[0].ns;
184 let speed = |m: &Measured| -> Vec<f64> { m.ns.iter().zip(base).map(|(a, b)| b / a).collect() };
185 let control = data.iter().find(|m| m.name.ends_with("(A/A control)"));
186 let noise = control.map_or(0.0, |c| {
187 let (lo, hi) = bootstrap(&speed(c));
188 (1.0 - lo).abs().max((hi - 1.0).abs())
189 });
190 let _ = writeln!(out, "\n{group} (noise floor ±{:.9}%)", noise * 100.0);
191 for m in data {
192 let med = median(&mut m.ns.clone());
193 let r = speed(m);
194 let (lo, hi) = bootstrap(&r);
195 let rm = median(&mut r.clone());
196 let verdict = if std::ptr::eq(m, &data[0]) {
197 "baseline".to_string()
198 } else if lo > 1.0 + noise {
199 format!("faster {rm:.9}x")
200 } else if hi < 1.0 - noise {
201 format!("slower {:.9}x", 1.0 / rm)
202 } else {
203 "inside noise".to_string()
204 };
205 let _ = writeln!(out, " {:<34} {:>14} speed {:.9} [{:.9}, {:.9}] {verdict}", m.name, fmt_rate(1e9 / med), rm, lo, hi);
206 }
207 out
208}
209
210fn fmt_rate(per_s: f64) -> String {
211 if per_s >= 1e9 {
212 format!("{:.9} billion/s", per_s / 1e9)
213 } else if per_s >= 1e6 {
214 format!("{:.9} million/s", per_s / 1e6)
215 } else if per_s >= 1e3 {
216 format!("{:.9} thousand/s", per_s / 1e3)
217 } else {
218 format!("{per_s:.9} per second")
219 }
220}
221
222fn pin_and_raise() {
223 #[cfg(windows)]
224 {
225 #[link(name = "kernel32")]
226 extern "system" {
227 fn GetCurrentThread() -> *mut core::ffi::c_void;
228 fn SetThreadAffinityMask(thread: *mut core::ffi::c_void, mask: usize) -> usize;
229 fn SetThreadPriority(thread: *mut core::ffi::c_void, priority: i32) -> i32;
230 }
231 const THREAD_PRIORITY_HIGHEST: i32 = 2;
232 unsafe {
233 let t = GetCurrentThread();
234 SetThreadAffinityMask(t, 1usize << 2);
235 SetThreadPriority(t, THREAD_PRIORITY_HIGHEST);
236 }
237 }
238}
239
240mod clock {
241 use std::sync::OnceLock;
242 use std::time::Instant;
243
244 fn origin() -> Instant {
245 static O: OnceLock<Instant> = OnceLock::new();
246 *O.get_or_init(Instant::now)
247 }
248
249 #[inline(always)]
250 pub fn now() -> u64 {
251 #[cfg(target_arch = "x86_64")]
252 {
253 unsafe { core::arch::x86_64::_rdtsc() }
254 }
255 #[cfg(not(target_arch = "x86_64"))]
256 {
257 origin().elapsed().as_nanos() as u64
258 }
259 }
260
261 fn ns_per_tick() -> f64 {
262 static K: OnceLock<f64> = OnceLock::new();
263 *K.get_or_init(|| {
264 if cfg!(not(target_arch = "x86_64")) {
265 return 1.0;
266 }
267 let (t0, c0) = (origin().elapsed(), now());
268 let spin = Instant::now();
269 while spin.elapsed().as_millis() < 50 {
270 std::hint::spin_loop();
271 }
272 let (t1, c1) = (origin().elapsed(), now());
273 (t1 - t0).as_nanos() as f64 / c1.wrapping_sub(c0).max(1) as f64
274 })
275 }
276
277 pub fn to_ns(ticks: u64) -> f64 {
278 ticks as f64 * ns_per_tick()
279 }
280}
281
282#[cfg(test)]
283mod tests {
284 use super::*;
285
286 #[test]
287 fn detects_a_real_difference_and_passes_the_control() {
288 let v: Vec<u64> = (0..4096).collect();
289 let mut arms: Vec<Arm<'_>> = vec![
290 ("sum".into(), Box::new(|b: &mut Bencher| b.iter(|| slain(&v).iter().sum::<u64>()))),
291 ("sum x4".into(), Box::new(|b: &mut Bencher| b.iter(|| (0..4).map(|_| slain(&v).iter().sum::<u64>()).sum::<u64>()))),
292 ];
293 let cfg = Config { warm: Duration::from_millis(20), sample: Duration::from_micros(200), rounds: 15 };
294 let out = run("test", &mut arms, cfg);
295 assert!(out.contains("slower"), "{out}");
296 assert!(out.contains("(A/A control)"), "{out}");
297 assert!(out.contains("/s"), "{out}");
298 }
299}