1#![allow(clippy::too_many_arguments)]
2
3use std::{borrow::Cow, cell::RefCell, fmt, num::NonZeroUsize, time::Duration};
4
5use clap::ColorChoice;
6use regex::Regex;
7
8use crate::{
9 bench::BenchOptions,
10 config::{Action, Filter, ParsedSeconds, RunIgnored, SortingAttr},
11 counter::{
12 BytesCount, BytesFormat, CharsCount, CyclesCount, IntoCounter, ItemsCount, MaxCountUInt,
13 PrivBytesFormat,
14 },
15 entry::{AnyBenchEntry, BenchEntryRunner, EntryTree},
16 thread_pool::BENCH_POOL,
17 time::{Timer, TimerKind},
18 tree_painter::{TreeColumn, TreePainter},
19 util::{self, defer},
20 Bencher,
21};
22
23#[derive(Default)]
25pub struct Divan {
26 action: Action,
27 timer: TimerKind,
28 reverse_sort: bool,
29 sorting_attr: SortingAttr,
30 color: ColorChoice,
31 bytes_format: BytesFormat,
32 filters: Vec<Filter>,
33 skip_filters: Vec<Filter>,
34 run_ignored: RunIgnored,
35 bench_options: BenchOptions<'static>,
36}
37
38pub(crate) struct SharedContext {
40 pub action: Action,
42
43 pub timer: Timer,
45}
46
47impl fmt::Debug for Divan {
48 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
49 f.debug_struct("Divan").finish_non_exhaustive()
50 }
51}
52
53impl Divan {
54 pub fn main(&self) {
58 self.run_action(self.action);
59 }
60
61 pub fn run_benches(&self) {
63 self.run_action(Action::Bench);
64 }
65
66 pub fn test_benches(&self) {
71 self.run_action(Action::Test);
72 }
73
74 pub fn list_benches(&self) {
76 self.run_action(Action::Test);
77 }
78
79 fn filter(&self, entry_path: &str) -> bool {
85 if !self.filters.is_empty()
86 && !self.filters.iter().any(|filter| filter.is_match(entry_path))
87 {
88 return false;
89 }
90
91 !self.skip_filters.iter().any(|filter| filter.is_match(entry_path))
92 }
93
94 pub(crate) fn should_ignore(&self, ignored: bool) -> bool {
95 !self.run_ignored.should_run(ignored)
96 }
97
98 pub(crate) fn run_action(&self, action: Action) {
99 let _drop_threads = defer(|| BENCH_POOL.drop_threads());
100
101 let mut tree: Vec<EntryTree> = if cfg!(miri) {
102 Vec::new()
104 } else {
105 let group_entries = &crate::entry::GROUP_ENTRIES;
106
107 let generic_bench_entries = group_entries
108 .iter()
109 .flat_map(|group| group.generic_benches_iter().map(AnyBenchEntry::GenericBench));
110
111 let bench_entries = crate::entry::BENCH_ENTRIES
112 .iter()
113 .map(AnyBenchEntry::Bench)
114 .chain(generic_bench_entries);
115
116 let mut tree = EntryTree::from_benches(bench_entries);
117
118 for group in group_entries.iter() {
119 EntryTree::insert_group(&mut tree, group);
120 }
121
122 tree
123 };
124
125 EntryTree::retain(&mut tree, |entry_path| self.filter(entry_path));
128
129 if tree.is_empty() {
131 return;
132 }
133
134 EntryTree::sort_by_attr(&mut tree, self.sorting_attr, self.reverse_sort);
136
137 let timer = match self.timer {
138 TimerKind::Os => Timer::Os,
139
140 TimerKind::Tsc => {
141 match Timer::get_tsc() {
142 Ok(tsc) => tsc,
143 Err(error) => {
144 eprintln!("warning: CPU timestamp counter is unavailable ({error}), defaulting to OS");
145 Timer::Os
146 }
147 }
148 }
149 };
150
151 if action.is_bench() {
152 eprintln!("Timer precision: {}", timer.precision());
153 }
154
155 let shared_context = SharedContext { action, timer };
156
157 let column_widths = if action.is_bench() {
158 TreeColumn::ALL.map(|column| {
159 if column.is_last() {
160 0
162 } else {
163 EntryTree::common_column_width(&tree, column)
164 }
165 })
166 } else {
167 [0; TreeColumn::COUNT]
168 };
169
170 let tree_painter =
171 RefCell::new(TreePainter::new(EntryTree::max_name_span(&tree, 0), column_widths));
172
173 self.run_tree(action, &tree, &shared_context, None, &tree_painter);
174 }
175
176 fn run_tree(
177 &self,
178 action: Action,
179 tree: &[EntryTree],
180 shared_context: &SharedContext,
181 parent_options: Option<&BenchOptions>,
182 tree_painter: &RefCell<TreePainter>,
183 ) {
184 for (i, child) in tree.iter().enumerate() {
185 let is_last = i == tree.len() - 1;
186
187 let name = child.display_name();
188
189 let child_options = child.bench_options();
190
191 let options: BenchOptions;
193 let options: Option<&BenchOptions> = match (parent_options, child_options) {
194 (None, None) => None,
195 (Some(options), None) | (None, Some(options)) => Some(options),
196 (Some(parent_options), Some(child_options)) => {
197 options = child_options.overwrite(parent_options);
198 Some(&options)
199 }
200 };
201
202 match child {
203 EntryTree::Leaf { entry, args } => self.run_bench_entry(
204 action,
205 *entry,
206 args.as_deref(),
207 shared_context,
208 options,
209 tree_painter,
210 is_last,
211 ),
212 EntryTree::Parent { children, .. } => {
213 tree_painter.borrow_mut().start_parent(name, is_last);
214
215 self.run_tree(action, children, shared_context, options, tree_painter);
216
217 tree_painter.borrow_mut().finish_parent();
218 }
219 }
220 }
221 }
222
223 fn run_bench_entry(
224 &self,
225 action: Action,
226 bench_entry: AnyBenchEntry,
227 bench_arg_names: Option<&[&&str]>,
228 shared_context: &SharedContext,
229 entry_options: Option<&BenchOptions>,
230 tree_painter: &RefCell<TreePainter>,
231 is_last_entry: bool,
232 ) {
233 use crate::bench::BenchContext;
234
235 let entry_display_name = bench_entry.display_name();
236
237 let options: BenchOptions;
239 let options: &BenchOptions = match entry_options {
240 None => &self.bench_options,
241 Some(entry_options) => {
242 options = self.bench_options.overwrite(entry_options);
243 &options
244 }
245 };
246
247 if self.should_ignore(options.ignore.unwrap_or_default()) {
248 tree_painter.borrow_mut().ignore_leaf(entry_display_name, is_last_entry);
249 return;
250 }
251
252 if action.is_list() {
254 let mut tree_painter = tree_painter.borrow_mut();
255 tree_painter.start_leaf(entry_display_name, is_last_entry);
256 tree_painter.finish_empty_leaf();
257 return;
258 }
259
260 let mut thread_counts: Vec<NonZeroUsize> = options
261 .threads
262 .as_deref()
263 .unwrap_or_default()
264 .iter()
265 .map(|&n| match NonZeroUsize::new(n) {
266 Some(n) => n,
267 None => crate::util::known_parallelism(),
268 })
269 .collect();
270
271 thread_counts.sort_unstable();
272 thread_counts.dedup();
273
274 let thread_counts: &[NonZeroUsize] =
275 if thread_counts.is_empty() { &[NonZeroUsize::MIN] } else { &thread_counts };
276
277 let has_thread_branches = thread_counts.len() > 1;
279
280 let run_bench = |bench_display_name: &str,
281 is_last_bench: bool,
282 with_bencher: &dyn Fn(Bencher)| {
283 if has_thread_branches {
284 tree_painter.borrow_mut().start_parent(bench_display_name, is_last_bench);
285 } else {
286 tree_painter.borrow_mut().start_leaf(bench_display_name, is_last_bench);
287 }
288
289 for (i, &thread_count) in thread_counts.iter().enumerate() {
290 let is_last_thread_count =
291 if has_thread_branches { i == thread_counts.len() - 1 } else { is_last_bench };
292
293 if has_thread_branches {
294 tree_painter
295 .borrow_mut()
296 .start_leaf(&format!("t={thread_count}"), is_last_thread_count);
297 }
298
299 let mut bench_context = BenchContext::new(shared_context, options, thread_count);
300 with_bencher(Bencher::new(&mut bench_context));
301
302 if !bench_context.did_run {
303 eprintln!(
304 "warning: No benchmark function registered for '{bench_display_name}'"
305 );
306 }
307
308 let should_compute_stats =
309 bench_context.did_run && shared_context.action.is_bench();
310
311 if should_compute_stats {
312 let stats = bench_context.compute_stats();
313 codspeed::collect_walltime_results(
314 &bench_context,
315 &bench_entry,
316 bench_display_name,
317 );
318 tree_painter.borrow_mut().finish_leaf(
319 is_last_thread_count,
320 &stats,
321 self.bytes_format,
322 );
323 } else {
324 tree_painter.borrow_mut().finish_empty_leaf();
325 }
326 }
327
328 if has_thread_branches {
329 tree_painter.borrow_mut().finish_parent();
330 }
331 };
332
333 match bench_entry.bench_runner() {
334 BenchEntryRunner::Plain(bench) => run_bench(entry_display_name, is_last_entry, bench),
335
336 BenchEntryRunner::Args(bench_runner) => {
337 tree_painter.borrow_mut().start_parent(entry_display_name, is_last_entry);
338
339 let bench_runner = bench_runner();
340 let orig_arg_names = bench_runner.arg_names();
341 let bench_arg_names = bench_arg_names.unwrap_or_default();
342
343 for (i, &arg_name) in bench_arg_names.iter().enumerate() {
344 let is_last_arg = i == bench_arg_names.len() - 1;
345 let arg_index = util::slice_ptr_index(orig_arg_names, arg_name);
346
347 run_bench(arg_name, is_last_arg, &|bencher| {
348 bench_runner.bench(bencher, arg_index);
349 });
350 }
351
352 tree_painter.borrow_mut().finish_parent();
353 }
354 }
355 }
356}
357
358mod codspeed {
359 use crate::bench::BenchContext;
360 use crate::entry::AnyBenchEntry;
361
362 pub(crate) fn collect_walltime_results(
363 bench_context: &BenchContext,
364 bench_entry: &AnyBenchEntry,
365 closure_bench_display_name: &str,
366 ) {
367 let (bench_name, uri) = {
370 let bench_function_name = bench_entry.meta().display_name;
371
372 let (bench_type_name, bench_arg_name) = {
373 let bench_function_or_type_name = bench_entry.display_name().to_string();
374
375 let type_name = if bench_function_or_type_name == bench_function_name {
376 None
377 } else {
378 Some(bench_function_or_type_name)
379 };
380
381 let arg_name = match type_name.as_ref() {
382 None => {
383 if closure_bench_display_name == bench_function_name {
384 None
385 } else {
386 Some(closure_bench_display_name)
387 }
388 }
389 Some(type_name) => {
390 if closure_bench_display_name == type_name {
391 None
392 } else {
393 Some(closure_bench_display_name)
394 }
395 }
396 };
397
398 (type_name, arg_name)
399 };
400
401 let mut bench_name = bench_function_name.to_string();
402
403 match (bench_type_name, bench_arg_name) {
404 (None, None) => {}
405 (Some(type_name), None) => {
406 bench_name.push_str(format!("[{type_name}]").as_str());
407 }
408 (None, Some(arg_name)) => {
409 bench_name.push_str(format!("[{arg_name}]").as_str());
410 }
411 (Some(type_name), Some(arg_name)) => {
412 bench_name.push_str(format!("[{type_name}, {arg_name}]").as_str());
413 }
414 }
415
416 let file = bench_entry.meta().location.file;
417 let mut module_path =
418 bench_entry.meta().module_path_components().skip(1).collect::<Vec<_>>().join("::");
419 if !module_path.is_empty() {
420 module_path.push_str("::");
421 }
422 let uri = format!("{file}::{module_path}{bench_name}");
423 (bench_name, uri)
424 };
425
426 let iter_per_round = bench_context.samples.sample_size;
427 let iters_per_round =
428 vec![iter_per_round as u128; bench_context.samples.time_samples.len()];
429 let times_per_round_ns: Vec<_> =
430 bench_context.samples.time_samples.iter().map(|s| s.duration.picos / 1_000).collect();
431 let max_time_ns = bench_context.options.max_time.map(|t| t.as_nanos());
432
433 if let Err(error) = ::codspeed::fifo::send_cmd(codspeed::fifo::Command::CurrentBenchmark {
434 pid: std::process::id(),
435 uri: uri.clone(),
436 }) {
437 if codspeed::utils::running_with_codspeed_runner() {
438 eprintln!("Failed to send benchmark URI to runner: {error:?}");
439 }
440 }
441
442 ::codspeed::walltime_results::WalltimeBenchmark::collect_raw_walltime_results(
443 "divan",
444 bench_name,
445 uri,
446 iters_per_round,
447 times_per_round_ns,
448 max_time_ns,
449 );
450 }
451}
452
453pub trait SkipRegex {
455 fn skip_regex(self, divan: &mut Divan);
456}
457
458impl SkipRegex for Regex {
459 fn skip_regex(self, divan: &mut Divan) {
460 divan.skip_filters.push(Filter::Regex(self));
461 }
462}
463
464impl SkipRegex for &str {
465 #[track_caller]
466 fn skip_regex(self, divan: &mut Divan) {
467 Regex::new(self).unwrap().skip_regex(divan);
468 }
469}
470
471impl SkipRegex for String {
472 #[track_caller]
473 fn skip_regex(self, divan: &mut Divan) {
474 self.as_str().skip_regex(divan)
475 }
476}
477
478impl Divan {
480 pub fn from_args() -> Self {
482 Self::default().config_with_args()
483 }
484
485 #[must_use]
489 pub fn config_with_args(mut self) -> Self {
490 let mut command = crate::cli::command();
491
492 let matches = command.get_matches_mut();
493 let is_exact = matches.get_flag("exact");
494
495 let mut parse_filter = |filter: &String| {
496 if is_exact {
497 Filter::Exact(filter.to_owned())
498 } else {
499 match Regex::new(filter) {
500 Ok(r) => Filter::Regex(r),
501 Err(error) => {
502 let kind = clap::error::ErrorKind::ValueValidation;
503 command.error(kind, error).exit();
504 }
505 }
506 }
507 };
508
509 if let Some(filters) = matches.get_many::<String>("filter") {
510 self.filters.extend(filters.map(&mut parse_filter));
511 }
512
513 if let Some(skip_filters) = matches.get_many::<String>("skip") {
514 self.skip_filters.extend(skip_filters.map(&mut parse_filter));
515 }
516
517 self.action = if matches.get_flag("list") {
518 Action::List
519 } else if matches.get_flag("test") || !matches.get_flag("bench") {
520 Action::Test
524 } else {
525 Action::Bench
526 };
527
528 if let Some(&color) = matches.get_one("color") {
529 self.color = color;
530 }
531
532 if matches.get_flag("ignored") {
533 self.run_ignored = RunIgnored::Only;
534 } else if matches.get_flag("include-ignored") {
535 self.run_ignored = RunIgnored::Yes;
536 }
537
538 if let Some(&timer) = matches.get_one("timer") {
539 self.timer = timer;
540 }
541
542 if let Some(&sorting_attr) = matches.get_one("sortr") {
543 self.reverse_sort = true;
544 self.sorting_attr = sorting_attr;
545 } else if let Some(&sorting_attr) = matches.get_one("sort") {
546 self.reverse_sort = false;
547 self.sorting_attr = sorting_attr;
548 }
549
550 if let Some(&sample_count) = matches.get_one("sample-count") {
551 self.bench_options.sample_count = Some(sample_count);
552 }
553
554 if let Some(&sample_size) = matches.get_one("sample-size") {
555 self.bench_options.sample_size = Some(sample_size);
556 }
557
558 if let Some(thread_counts) = matches.get_many::<usize>("threads") {
559 let mut threads: Vec<usize> = thread_counts.copied().collect();
560 threads.sort_unstable();
561 threads.dedup();
562 self.bench_options.threads = Some(Cow::Owned(threads));
563 }
564
565 if let Some(&ParsedSeconds(min_time)) = matches.get_one("min-time") {
566 self.bench_options.min_time = Some(min_time);
567 }
568
569 if let Some(&ParsedSeconds(max_time)) = matches.get_one("max-time") {
570 self.bench_options.max_time = Some(max_time);
571 }
572
573 if let Some(mut skip_ext_time) = matches.get_many::<bool>("skip-ext-time") {
574 self.bench_options.skip_ext_time =
576 Some(matches!(skip_ext_time.next(), Some(true) | None));
577 }
578
579 if let Some(&count) = matches.get_one::<MaxCountUInt>("items-count") {
580 self.counter_mut(ItemsCount::new(count));
581 }
582
583 if let Some(&count) = matches.get_one::<MaxCountUInt>("bytes-count") {
584 self.counter_mut(BytesCount::new(count));
585 }
586
587 if let Some(&PrivBytesFormat(bytes_format)) = matches.get_one("bytes-format") {
588 self.bytes_format = bytes_format;
589 }
590
591 if let Some(&count) = matches.get_one::<MaxCountUInt>("chars-count") {
592 self.counter_mut(CharsCount::new(count));
593 }
594
595 if let Some(&count) = matches.get_one::<MaxCountUInt>("cycles-count") {
596 self.counter_mut(CyclesCount::new(count));
597 }
598
599 self
600 }
601
602 #[must_use]
607 pub fn color(mut self, yes: impl Into<Option<bool>>) -> Self {
608 self.color = match yes.into() {
609 None => ColorChoice::Auto,
610 Some(true) => ColorChoice::Always,
611 Some(false) => ColorChoice::Never,
612 };
613 self
614 }
615
616 #[must_use]
620 pub fn run_ignored(mut self) -> Self {
621 self.run_ignored = RunIgnored::Yes;
622 self
623 }
624
625 #[must_use]
629 pub fn run_only_ignored(mut self) -> Self {
630 self.run_ignored = RunIgnored::Only;
631 self
632 }
633
634 #[must_use]
670 pub fn skip_regex(mut self, filter: impl SkipRegex) -> Self {
671 filter.skip_regex(&mut self);
672 self
673 }
674
675 #[must_use]
698 pub fn skip_exact(mut self, filter: impl Into<String>) -> Self {
699 self.skip_filters.push(Filter::Exact(filter.into()));
700 self
701 }
702
703 #[inline]
712 pub fn sample_count(mut self, count: u32) -> Self {
713 self.bench_options.sample_count = Some(count);
714 self
715 }
716
717 #[inline]
721 pub fn sample_size(mut self, count: u32) -> Self {
722 self.bench_options.sample_size = Some(count);
723 self
724 }
725
726 #[inline]
735 pub fn threads<T>(mut self, threads: T) -> Self
736 where
737 T: IntoIterator<Item = usize>,
738 {
739 self.bench_options.threads = {
740 let mut threads: Vec<usize> = threads.into_iter().collect();
741 threads.sort_unstable();
742 threads.dedup();
743 Some(Cow::Owned(threads))
744 };
745 self
746 }
747
748 #[inline]
752 pub fn min_time(mut self, time: Duration) -> Self {
753 self.bench_options.min_time = Some(time);
754 self
755 }
756
757 #[inline]
761 pub fn max_time(mut self, time: Duration) -> Self {
762 self.bench_options.max_time = Some(time);
763 self
764 }
765
766 #[inline]
771 pub fn skip_ext_time(mut self, skip: bool) -> Self {
772 self.bench_options.skip_ext_time = Some(skip);
773 self
774 }
775}
776
777impl Divan {
780 #[inline]
781 fn counter_mut<C: IntoCounter>(&mut self, counter: C) -> &mut Self {
782 self.bench_options.counters.insert(counter);
783 self
784 }
785
786 #[inline]
788 pub fn counter<C: IntoCounter>(mut self, counter: C) -> Self {
789 self.counter_mut(counter);
790 self
791 }
792
793 #[inline]
798 pub fn items_count<C: Into<ItemsCount>>(self, count: C) -> Self {
799 self.counter(count.into())
800 }
801
802 #[inline]
807 pub fn bytes_count<C: Into<BytesCount>>(self, count: C) -> Self {
808 self.counter(count.into())
809 }
810
811 #[inline]
816 pub fn bytes_format(mut self, format: BytesFormat) -> Self {
817 self.bytes_format = format;
818 self
819 }
820
821 #[inline]
826 pub fn chars_count<C: Into<CharsCount>>(self, count: C) -> Self {
827 self.counter(count.into())
828 }
829
830 #[inline]
835 pub fn cycles_count<C: Into<CyclesCount>>(self, count: C) -> Self {
836 self.counter(count.into())
837 }
838}