1use std::{cmp::Ordering, ptr::NonNull};
2
3use crate::{
4 benchmark::{BenchOptions, DEFAULT_SAMPLE_COUNT},
5 config::SortingAttr,
6 counter::KnownCounterKind,
7 entry::{AnyBenchEntry, EntryLocation, EntryMeta, GenericBenchEntry, GroupEntry},
8 tree_painter::TreeColumn,
9 util::sort::natural_cmp,
10};
11
12/// `BenchEntry` tree organized by path components.
13pub(crate) enum EntryTree<'a> {
14 /// Benchmark group; parent to leaves and other parents.
15 Parent { raw_name: &'a str, group: Option<&'a GroupEntry>, children: Vec<Self> },
16
17 /// Benchmark entry leaf.
18 Leaf {
19 /// The benchmark entry being run.
20 entry: AnyBenchEntry<'a>,
21
22 /// The names of arguments to run.
23 args: Option<Vec<&'static &'static str>>,
24 },
25}
26
27impl<'a> EntryTree<'a> {
28 /// Constructs a tree from an iterator of benchmark entries in the order
29 /// they're produced.
30 pub fn from_benches<I>(benches: I) -> Vec<Self>
31 where
32 I: IntoIterator<Item = AnyBenchEntry<'a>>,
33 {
34 let mut result = Vec::<Self>::new();
35
36 for bench in benches {
37 let mut insert_entry = |path_iter| {
38 Self::insert_entry(&mut result, bench, path_iter);
39 };
40
41 match bench {
42 AnyBenchEntry::Bench(bench) => {
43 insert_entry(&mut bench.meta.module_path_components());
44 }
45 AnyBenchEntry::GenericBench(bench) => {
46 insert_entry(&mut bench.path_components());
47 }
48 }
49 }
50
51 result
52 }
53
54 /// Returns the maximum span for a name in `tree`.
55 ///
56 /// This is the number of terminal columns used for labeling benchmark names
57 /// prior to emitting stats columns.
58 pub fn max_name_span(tree: &[Self], depth: usize) -> usize {
59 // The number of terminal columns used per-depth for box drawing
60 // characters. For example, "│ ╰─ " is 6 for depth 2.
61 const DEPTH_COLS: usize = 3;
62
63 tree.iter()
64 .map(|node| {
65 let node_name_span = {
66 let prefix_len = depth * DEPTH_COLS;
67 let name_len = node.display_name().chars().count();
68 prefix_len + name_len
69 };
70
71 // The maximum span of any descendent.
72 let children_max_span = Self::max_name_span(node.children(), depth + 1);
73
74 // The maximum span of any runtime argument.
75 let args_max_span = node
76 .arg_names()
77 .unwrap_or_default()
78 .iter()
79 .map(|arg| {
80 let prefix_len = (depth + 1) * DEPTH_COLS;
81 let name_len = arg.chars().count();
82 prefix_len + name_len
83 })
84 .max()
85 .unwrap_or_default();
86
87 node_name_span.max(children_max_span).max(args_max_span)
88 })
89 .max()
90 .unwrap_or_default()
91 }
92
93 /// Returns the likely span for a given column.
94 pub fn common_column_width(tree: &[Self], column: TreeColumn) -> usize {
95 // Time and throughput info.
96 if column.is_time_stat() {
97 return KnownCounterKind::MAX_COMMON_COLUMN_WIDTH;
98 }
99
100 tree.iter()
101 .map(|tree| {
102 let Some(options) = tree.bench_options() else {
103 return 0;
104 };
105
106 let width = match column {
107 TreeColumn::Samples => {
108 let sample_count = options.sample_count.unwrap_or(DEFAULT_SAMPLE_COUNT);
109 1 + sample_count.checked_ilog10().unwrap_or_default() as usize
110 }
111
112 // Iters is the last column, so it does not need pad width.
113 // All other columns are time stats handled previously.
114 _ => 0,
115 };
116
117 width.max(Self::common_column_width(tree.children(), column))
118 })
119 .max()
120 .unwrap_or_default()
121 }
122
123 /// Inserts the benchmark group into a tree.
124 ///
125 /// Groups are inserted after tree construction because it prevents having
126 /// parents without terminating leaves. Groups that do not match an existing
127 /// parent are not inserted.
128 pub fn insert_group(mut tree: &mut [Self], group: &'a GroupEntry) {
129 // Update `tree` to be the innermost set of subtrees whose parents match
130 // `group.module_path`.
131 'component: for component in group.meta.module_path_components() {
132 for subtree in tree {
133 match subtree {
134 EntryTree::Parent { raw_name, children, .. } if component == *raw_name => {
135 tree = children;
136 continue 'component;
137 }
138 _ => {}
139 }
140 }
141
142 // No matches for this component in any subtrees.
143 return;
144 }
145
146 // Find the matching tree to insert the group into.
147 for subtree in tree {
148 match subtree {
149 EntryTree::Parent { raw_name, group: slot, .. }
150 if group.meta.raw_name == *raw_name =>
151 {
152 *slot = Some(group);
153 return;
154 }
155 _ => {}
156 }
157 }
158 }
159
160 /// Removes entries from the tree whose paths do not match the filter.
161 pub fn retain(tree: &mut Vec<Self>, mut filter: impl FnMut(&str) -> bool) {
162 fn retain(
163 tree: &mut Vec<EntryTree>,
164 parent_path: &str,
165 filter: &mut impl FnMut(&str) -> bool,
166 ) {
167 tree.retain_mut(|subtree| {
168 let subtree_path: String;
169 let subtree_path: &str = if parent_path.is_empty() {
170 subtree.display_name()
171 } else {
172 subtree_path = format!("{parent_path}::{}", subtree.display_name());
173 &subtree_path
174 };
175
176 match subtree {
177 EntryTree::Parent { children, .. } => {
178 retain(children, subtree_path, filter);
179
180 // If no children exist, filter out this parent.
181 !children.is_empty()
182 }
183
184 EntryTree::Leaf { args: None, .. } => filter(subtree_path),
185
186 EntryTree::Leaf { args: Some(args), .. } => {
187 args.retain(|arg| filter(&format!("{subtree_path}::{arg}")));
188
189 // If no arguments exist, filter out this leaf.
190 !args.is_empty()
191 }
192 }
193 });
194 }
195 retain(tree, "", &mut filter);
196 }
197
198 /// Sorts the tree by the given ordering.
199 pub fn sort_by_attr(tree: &mut [Self], attr: SortingAttr, reverse: bool) {
200 let apply_reverse =
201 |ordering: Ordering| if reverse { ordering.reverse() } else { ordering };
202
203 tree.sort_unstable_by(|a, b| apply_reverse(a.cmp_by_attr(b, attr)));
204
205 tree.iter_mut().for_each(|tree| {
206 match tree {
207 // Sort benchmark arguments.
208 EntryTree::Leaf { args, .. } => {
209 if let Some(args) = args {
210 args.sort_by(|&a, &b| apply_reverse(attr.cmp_bench_arg_names(a, b)));
211 }
212 }
213
214 // Sort children.
215 EntryTree::Parent { children, .. } => {
216 Self::sort_by_attr(children, attr, reverse);
217 }
218 }
219 });
220 }
221
222 fn cmp_by_attr(&self, other: &Self, attr: SortingAttr) -> Ordering {
223 // We take advantage of the fact that entries have stable addresses,
224 // unlike `EntryTree`.
225 let entry_addr_ordering = match (self.entry_addr(), other.entry_addr()) {
226 (Some(a), Some(b)) => Some(a.cmp(&b)),
227 _ => None,
228 };
229
230 // If entries have the same address, then all attributes will be equal.
231 if matches!(entry_addr_ordering, Some(Ordering::Equal)) {
232 return Ordering::Equal;
233 }
234
235 for attr in attr.with_tie_breakers() {
236 let ordering = match attr {
237 SortingAttr::Kind => self.kind().cmp(&other.kind()),
238 SortingAttr::Name => self.cmp_display_name(other),
239 SortingAttr::Location => {
240 let location_ordering = self.location().cmp(&other.location());
241
242 // Use the entry's address to break location ties.
243 //
244 // This makes generic benchmarks use the same order as their
245 // types and constants.
246 if location_ordering.is_eq() {
247 entry_addr_ordering.unwrap_or(Ordering::Equal)
248 } else {
249 location_ordering
250 }
251 }
252 };
253
254 if ordering.is_ne() {
255 return ordering;
256 }
257 }
258
259 Ordering::Equal
260 }
261
262 /// Helper for constructing a tree.
263 ///
264 /// This uses recursion because the iterative approach runs into limitations
265 /// with mutable borrows.
266 fn insert_entry(
267 tree: &mut Vec<Self>,
268 entry: AnyBenchEntry<'a>,
269 rem_modules: &mut dyn Iterator<Item = &'a str>,
270 ) {
271 let Some(current_module) = rem_modules.next() else {
272 tree.push(Self::Leaf {
273 entry,
274 args: entry.arg_names().map(|args| args.iter().collect()),
275 });
276 return;
277 };
278
279 let Some(children) = Self::get_children(tree, current_module) else {
280 tree.push(Self::from_path(entry, current_module, rem_modules));
281 return;
282 };
283
284 Self::insert_entry(children, entry, rem_modules);
285 }
286
287 /// Constructs a sequence of branches from a module path.
288 fn from_path(
289 entry: AnyBenchEntry<'a>,
290 current_module: &'a str,
291 rem_modules: &mut dyn Iterator<Item = &'a str>,
292 ) -> Self {
293 let child = if let Some(next_module) = rem_modules.next() {
294 Self::from_path(entry, next_module, rem_modules)
295 } else {
296 Self::Leaf { entry, args: entry.arg_names().map(|args| args.iter().collect()) }
297 };
298 Self::Parent { raw_name: current_module, group: None, children: vec![child] }
299 }
300
301 /// Finds the `Parent.children` for the corresponding module in `tree`.
302 fn get_children<'t>(tree: &'t mut [Self], module: &str) -> Option<&'t mut Vec<Self>> {
303 tree.iter_mut().find_map(|tree| match tree {
304 Self::Parent { raw_name, children, group: _ } if *raw_name == module => Some(children),
305 _ => None,
306 })
307 }
308
309 /// Returns an integer denoting the enum variant.
310 ///
311 /// This is used instead of `std::mem::Discriminant` because it does not
312 /// implement `Ord`.
313 pub fn kind(&self) -> i32 {
314 // Leaves should appear before parents.
315 match self {
316 Self::Leaf { .. } => 0,
317 Self::Parent { .. } => 1,
318 }
319 }
320
321 /// Returns a pointer to use as the identity of the entry.
322 pub fn entry_addr(&self) -> Option<NonNull<()>> {
323 match self {
324 Self::Leaf { entry, .. } => Some(entry.entry_addr()),
325 Self::Parent { group, .. } => {
326 group.map(|entry: &GroupEntry| NonNull::from(entry).cast())
327 }
328 }
329 }
330
331 pub fn meta(&self) -> Option<&'a EntryMeta> {
332 match self {
333 Self::Parent { group, .. } => Some(&(*group)?.meta),
334 Self::Leaf { entry, .. } => Some(entry.meta()),
335 }
336 }
337
338 pub fn bench_options(&self) -> Option<&'a BenchOptions> {
339 self.meta()?.bench_options()
340 }
341
342 pub fn raw_name(&self) -> &'a str {
343 match self {
344 Self::Parent { group: Some(group), .. } => group.meta.raw_name,
345 Self::Parent { raw_name, .. } => raw_name,
346 Self::Leaf { entry, .. } => entry.raw_name(),
347 }
348 }
349
350 pub fn display_name(&self) -> &'a str {
351 if let Self::Leaf { entry, .. } = self {
352 entry.display_name()
353 } else if let Some(common) = self.meta() {
354 common.display_name
355 } else {
356 let raw_name = self.raw_name();
357 raw_name.strip_prefix("r#").unwrap_or(raw_name)
358 }
359 }
360
361 /// Returns the location of this entry, group, or the children's earliest
362 /// location.
363 fn location(&self) -> Option<&'a EntryLocation> {
364 if let Some(common) = self.meta() {
365 Some(&common.location)
366 } else {
367 self.children().iter().flat_map(Self::location).min()
368 }
369 }
370
371 /// Compares display names naturally, taking into account integers.
372 ///
373 /// There is special consideration for the `PartialOrd` implementation of
374 /// constants, so that `EntryConst` can sort integers and floats by value
375 /// instead of lexicographically.
376 fn cmp_display_name(&self, other: &Self) -> Ordering {
377 match (self, other) {
378 (
379 Self::Leaf {
380 entry:
381 AnyBenchEntry::GenericBench(GenericBenchEntry {
382 const_value: Some(this), ..
383 }),
384 ..
385 },
386 Self::Leaf {
387 entry:
388 AnyBenchEntry::GenericBench(GenericBenchEntry {
389 const_value: Some(other), ..
390 }),
391 ..
392 },
393 ) => this.cmp_name(other),
394
395 _ => natural_cmp(self.display_name(), other.display_name()),
396 }
397 }
398
399 fn children(&self) -> &[Self] {
400 match self {
401 Self::Leaf { .. } => &[],
402 Self::Parent { children, .. } => children,
403 }
404 }
405
406 fn arg_names(&self) -> Option<&[&'static &'static str]> {
407 match self {
408 Self::Leaf { args, .. } => args.as_deref(),
409 Self::Parent { .. } => None,
410 }
411 }
412}
413