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 #[cfg(unix)]
434 {
435 if let Err(error) =
436 ::codspeed::fifo::send_cmd(codspeed::fifo::Command::CurrentBenchmark {
437 pid: std::process::id(),
438 uri: uri.clone(),
439 })
440 {
441 if codspeed::utils::running_with_codspeed_runner() {
442 eprintln!("Failed to send benchmark URI to runner: {error:?}");
443 }
444 }
445 }
446
447 ::codspeed::walltime_results::WalltimeBenchmark::collect_raw_walltime_results(
448 "divan",
449 bench_name,
450 uri,
451 iters_per_round,
452 times_per_round_ns,
453 max_time_ns,
454 );
455 }
456}
457
458pub trait SkipRegex {
460 fn skip_regex(self, divan: &mut Divan);
461}
462
463impl SkipRegex for Regex {
464 fn skip_regex(self, divan: &mut Divan) {
465 divan.skip_filters.push(Filter::Regex(self));
466 }
467}
468
469impl SkipRegex for &str {
470 #[track_caller]
471 fn skip_regex(self, divan: &mut Divan) {
472 Regex::new(self).unwrap().skip_regex(divan);
473 }
474}
475
476impl SkipRegex for String {
477 #[track_caller]
478 fn skip_regex(self, divan: &mut Divan) {
479 self.as_str().skip_regex(divan)
480 }
481}
482
483impl Divan {
485 pub fn from_args() -> Self {
487 Self::default().config_with_args()
488 }
489
490 #[must_use]
494 pub fn config_with_args(mut self) -> Self {
495 let mut command = crate::cli::command();
496
497 let matches = command.get_matches_mut();
498 let is_exact = matches.get_flag("exact");
499
500 let mut parse_filter = |filter: &String| {
501 if is_exact {
502 Filter::Exact(filter.to_owned())
503 } else {
504 match Regex::new(filter) {
505 Ok(r) => Filter::Regex(r),
506 Err(error) => {
507 let kind = clap::error::ErrorKind::ValueValidation;
508 command.error(kind, error).exit();
509 }
510 }
511 }
512 };
513
514 if let Some(filters) = matches.get_many::<String>("filter") {
515 self.filters.extend(filters.map(&mut parse_filter));
516 }
517
518 if let Some(skip_filters) = matches.get_many::<String>("skip") {
519 self.skip_filters.extend(skip_filters.map(&mut parse_filter));
520 }
521
522 self.action = if matches.get_flag("list") {
523 Action::List
524 } else if matches.get_flag("test") || !matches.get_flag("bench") {
525 Action::Test
529 } else {
530 Action::Bench
531 };
532
533 if let Some(&color) = matches.get_one("color") {
534 self.color = color;
535 }
536
537 if matches.get_flag("ignored") {
538 self.run_ignored = RunIgnored::Only;
539 } else if matches.get_flag("include-ignored") {
540 self.run_ignored = RunIgnored::Yes;
541 }
542
543 if let Some(&timer) = matches.get_one("timer") {
544 self.timer = timer;
545 }
546
547 if let Some(&sorting_attr) = matches.get_one("sortr") {
548 self.reverse_sort = true;
549 self.sorting_attr = sorting_attr;
550 } else if let Some(&sorting_attr) = matches.get_one("sort") {
551 self.reverse_sort = false;
552 self.sorting_attr = sorting_attr;
553 }
554
555 if let Some(&sample_count) = matches.get_one("sample-count") {
556 self.bench_options.sample_count = Some(sample_count);
557 }
558
559 if let Some(&sample_size) = matches.get_one("sample-size") {
560 self.bench_options.sample_size = Some(sample_size);
561 }
562
563 if let Some(thread_counts) = matches.get_many::<usize>("threads") {
564 let mut threads: Vec<usize> = thread_counts.copied().collect();
565 threads.sort_unstable();
566 threads.dedup();
567 self.bench_options.threads = Some(Cow::Owned(threads));
568 }
569
570 if let Some(&ParsedSeconds(min_time)) = matches.get_one("min-time") {
571 self.bench_options.min_time = Some(min_time);
572 }
573
574 if let Some(&ParsedSeconds(max_time)) = matches.get_one("max-time") {
575 self.bench_options.max_time = Some(max_time);
576 }
577
578 if let Some(mut skip_ext_time) = matches.get_many::<bool>("skip-ext-time") {
579 self.bench_options.skip_ext_time =
581 Some(matches!(skip_ext_time.next(), Some(true) | None));
582 }
583
584 if let Some(&count) = matches.get_one::<MaxCountUInt>("items-count") {
585 self.counter_mut(ItemsCount::new(count));
586 }
587
588 if let Some(&count) = matches.get_one::<MaxCountUInt>("bytes-count") {
589 self.counter_mut(BytesCount::new(count));
590 }
591
592 if let Some(&PrivBytesFormat(bytes_format)) = matches.get_one("bytes-format") {
593 self.bytes_format = bytes_format;
594 }
595
596 if let Some(&count) = matches.get_one::<MaxCountUInt>("chars-count") {
597 self.counter_mut(CharsCount::new(count));
598 }
599
600 if let Some(&count) = matches.get_one::<MaxCountUInt>("cycles-count") {
601 self.counter_mut(CyclesCount::new(count));
602 }
603
604 self
605 }
606
607 #[must_use]
612 pub fn color(mut self, yes: impl Into<Option<bool>>) -> Self {
613 self.color = match yes.into() {
614 None => ColorChoice::Auto,
615 Some(true) => ColorChoice::Always,
616 Some(false) => ColorChoice::Never,
617 };
618 self
619 }
620
621 #[must_use]
625 pub fn run_ignored(mut self) -> Self {
626 self.run_ignored = RunIgnored::Yes;
627 self
628 }
629
630 #[must_use]
634 pub fn run_only_ignored(mut self) -> Self {
635 self.run_ignored = RunIgnored::Only;
636 self
637 }
638
639 #[must_use]
675 pub fn skip_regex(mut self, filter: impl SkipRegex) -> Self {
676 filter.skip_regex(&mut self);
677 self
678 }
679
680 #[must_use]
703 pub fn skip_exact(mut self, filter: impl Into<String>) -> Self {
704 self.skip_filters.push(Filter::Exact(filter.into()));
705 self
706 }
707
708 #[inline]
717 pub fn sample_count(mut self, count: u32) -> Self {
718 self.bench_options.sample_count = Some(count);
719 self
720 }
721
722 #[inline]
726 pub fn sample_size(mut self, count: u32) -> Self {
727 self.bench_options.sample_size = Some(count);
728 self
729 }
730
731 #[inline]
740 pub fn threads<T>(mut self, threads: T) -> Self
741 where
742 T: IntoIterator<Item = usize>,
743 {
744 self.bench_options.threads = {
745 let mut threads: Vec<usize> = threads.into_iter().collect();
746 threads.sort_unstable();
747 threads.dedup();
748 Some(Cow::Owned(threads))
749 };
750 self
751 }
752
753 #[inline]
757 pub fn min_time(mut self, time: Duration) -> Self {
758 self.bench_options.min_time = Some(time);
759 self
760 }
761
762 #[inline]
766 pub fn max_time(mut self, time: Duration) -> Self {
767 self.bench_options.max_time = Some(time);
768 self
769 }
770
771 #[inline]
776 pub fn skip_ext_time(mut self, skip: bool) -> Self {
777 self.bench_options.skip_ext_time = Some(skip);
778 self
779 }
780}
781
782impl Divan {
785 #[inline]
786 fn counter_mut<C: IntoCounter>(&mut self, counter: C) -> &mut Self {
787 self.bench_options.counters.insert(counter);
788 self
789 }
790
791 #[inline]
793 pub fn counter<C: IntoCounter>(mut self, counter: C) -> Self {
794 self.counter_mut(counter);
795 self
796 }
797
798 #[inline]
803 pub fn items_count<C: Into<ItemsCount>>(self, count: C) -> Self {
804 self.counter(count.into())
805 }
806
807 #[inline]
812 pub fn bytes_count<C: Into<BytesCount>>(self, count: C) -> Self {
813 self.counter(count.into())
814 }
815
816 #[inline]
821 pub fn bytes_format(mut self, format: BytesFormat) -> Self {
822 self.bytes_format = format;
823 self
824 }
825
826 #[inline]
831 pub fn chars_count<C: Into<CharsCount>>(self, count: C) -> Self {
832 self.counter(count.into())
833 }
834
835 #[inline]
840 pub fn cycles_count<C: Into<CyclesCount>>(self, count: C) -> Self {
841 self.counter(count.into())
842 }
843}