1use std::future::Future;
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::time::Instant;
4
5use crate::{Stats, Step};
6use crate::fmt_thousands_sep;
7use crate::timing_future::TimingFuture;
8#[cfg(feature = "track-allocator")]
9use crate::track_allocator::GLOBAL;
10
11pub struct Bencher {
12 pub name: String,
13 pub count: usize,
14 pub steps: Vec<Step>,
15 pub bytes: usize,
16 pub n: usize,
17 pub poll: usize,
18 pub format_fn: fn(&Stats, &Bencher),
19
20 pub mem_track: (&'static AtomicUsize, &'static AtomicUsize)
21}
22
23impl Bencher {
24 #[cfg(feature = "track-allocator")]
25 pub fn new(name: impl AsRef<str>, count: usize, bytes: usize) -> Self {
26 Bencher {
27 name: name.as_ref().to_owned(),
28 count,
29 steps: Vec::with_capacity(count),
30 bytes,
31 n: 0,
32 poll: 0,
33 format_fn: |s, b| Self::default_format(s, b),
34
35 mem_track: (GLOBAL.counter(), GLOBAL.peak())
36 }
37 }
38
39 #[cfg(not(feature = "track-allocator"))]
40 pub fn new(name: impl AsRef<str>, count: usize, bytes: usize, counter: &'static AtomicUsize, peak: &'static AtomicUsize) -> Self {
41 Bencher {
42 name: name.as_ref().to_owned(),
43 count,
44 steps: Vec::with_capacity(count),
45 bytes,
46 n: 0,
47 poll: 0,
48 format_fn: |s, b| Self::default_format(s, b),
49
50 mem_track: (counter, peak)
51 }
52 }
53
54 pub fn bench_once<T>(&self, f: &mut impl FnMut() -> T, n: usize) -> (u128, usize) {
56 let now = Instant::now();
57 self.reset_mem();
58
59 for _ in 0..n {
60 let _output = f();
61 }
62
63 (now.elapsed().as_nanos(), self.get_mem_peak())
64 }
65
66 pub fn iter<T>(&mut self, mut f: impl FnMut() -> T) {
67 let single = self.bench_once(&mut f, 1).0;
68 self.n = (1_000_000 / single.max(1)).max(1) as usize;
70 (0..self.count).for_each(|_| {
71 let res = self.bench_once(&mut f, self.n);
72 self.steps.push(Step {
73 time: res.0 / self.n as u128,
74 mem: res.1 / self.n
75 })
76 });
77 }
78
79 pub fn async_iter<'a, T, Fut: Future<Output=T>>(&'a mut self, mut f: impl FnMut() -> Fut + 'a) -> impl Future + 'a {
80 async move {
81 let single = TimingFuture::new(f()).await.elapsed_time.as_nanos();
82 self.n = (1_000_000 / single.max(1)).max(1) as usize;
84
85 let mut polls = Vec::with_capacity(self.count);
86
87 for _ in 0..self.count {
88 let mut mtime = 0u128;
89 self.reset_mem();
90
91 for _ in 0..self.n {
92 let tf = TimingFuture::new(f()).await;
93 mtime += tf.elapsed_time.as_nanos();
94 polls.push(tf.poll);
95 }
96
97 self.steps.push(Step {
98 time: mtime / self.n as u128,
99 mem: self.get_mem_peak() / self.n
100 });
101 }
102
103 self.poll = polls.iter().sum::<usize>() / polls.len();
104 }
105 }
106
107 pub fn finish(&self) {
108 let stats = Stats::from(&self.steps);
109 (self.format_fn)(&stats, self)
110 }
111
112 pub fn reset_mem(&self) {
113 self.mem_track.0.store(0, Ordering::SeqCst);
114 self.mem_track.1.store(0, Ordering::SeqCst);
115 }
116
117 pub fn get_mem_peak(&self) -> usize {
118 self.mem_track.1.load(Ordering::SeqCst)
119 }
120
121 fn default_format(stats: &Stats, bencher: &Bencher) {
122 bunt::println!(
123 "{[bg:white+blue+bold]} ... {[green+underline]} ns/iter (+/- {[red+underline]}) = {[yellow+underline]:.2} MB/s\
124 \n\t memory usage: {[green+underline]} bytes/iter (+/- {[red+underline]})\
125 \n\t @Total: {[magenta]} * {[white]} iters\
126 {[bold]}",
127 &bencher.name,
128 fmt_thousands_sep(stats.times_average, ','),
129 fmt_thousands_sep(stats.times_max - stats.times_min, ','),
130 (bencher.bytes as f64 * (1_000_000_000f64 / stats.times_average as f64)) / 1000f64 / 1000f64,
131
132 fmt_thousands_sep(stats.mem_average, ','),
133 fmt_thousands_sep(stats.mem_max - stats.mem_min, ','),
134
135 bencher.count,
136 bencher.n,
137
138 if bencher.poll > 0 {
139 format!(
140 "\n\t @avg {} polls ",
141 bencher.poll
142 )
143 } else {
144 String::new()
145 },
146 );
147 }
148}