1//! Happy little trees.
2
3use std::{io::Write, iter::repeat};
4
5use crate::{
6 alloc::{AllocOp, AllocTally},
7 counter::{AnyCounter, BytesFormat, KnownCounterKind},
8 stats::{Stats, StatsSet},
9 util,
10};
11
12const TREE_COL_BUF: usize = 2;
13
14/// Paints tree-style output using box-drawing characters.
15pub(crate) struct TreePainter {
16 /// The maximum number of characters taken by a name and its prefix. Emitted
17 /// information should be left-padded to start at this column.
18 max_name_span: usize,
19
20 column_widths: [usize; TreeColumn::COUNT],
21
22 depth: usize,
23
24 /// The current prefix to the name and content, e.g.
25 /// <code>│ │ </code> for three levels of nesting with the second level
26 /// being on the last node.
27 current_prefix: String,
28
29 /// Buffer for writing to before printing to stdout.
30 write_buf: String,
31}
32
33impl TreePainter {
34 pub fn new(max_name_span: usize, column_widths: [usize; TreeColumn::COUNT]) -> Self {
35 Self {
36 max_name_span,
37 column_widths,
38 depth: 0,
39 current_prefix: String::new(),
40 write_buf: String::new(),
41 }
42 }
43}
44
45impl TreePainter {
46 /// Enter a parent node.
47 pub fn start_parent(&mut self, name: &str, is_last: bool) {
48 let is_top_level = self.depth == 0;
49 let has_columns = self.has_columns();
50
51 let buf = &mut self.write_buf;
52 buf.clear();
53
54 let branch = if is_top_level {
55 ""
56 } else if !is_last {
57 "├─ "
58 } else {
59 "╰─ "
60 };
61 buf.extend([self.current_prefix.as_str(), branch, name]);
62
63 // Right-pad name if `has_columns`
64 if has_columns {
65 let max_span = self.max_name_span;
66 let buf_len = buf.chars().count();
67 let pad_len = TREE_COL_BUF + max_span.saturating_sub(buf_len);
68 buf.extend(repeat(' ').take(pad_len));
69
70 if buf_len > max_span {
71 self.max_name_span = buf_len;
72 }
73 }
74
75 // Write column headings.
76 if has_columns && is_top_level {
77 let names = TreeColumnData::from_fn(TreeColumn::name);
78 names.write(buf, &mut self.column_widths);
79 }
80
81 // Write column spacers.
82 if has_columns && !is_top_level {
83 TreeColumnData([""; TreeColumn::COUNT]).write(buf, &mut self.column_widths);
84 }
85
86 println!("{buf}");
87
88 self.depth += 1;
89
90 if !is_top_level {
91 self.current_prefix.push_str(if !is_last { "│ " } else { " " });
92 }
93 }
94
95 /// Exit the current parent node.
96 pub fn finish_parent(&mut self) {
97 self.depth -= 1;
98
99 // Improve legibility for multiple top-level parents.
100 if self.depth == 0 {
101 println!();
102 }
103
104 // The prefix is extended by 3 `char`s at a time.
105 let new_prefix_len = {
106 let mut iter = self.current_prefix.chars();
107 _ = iter.by_ref().rev().nth(2);
108 iter.as_str().len()
109 };
110 self.current_prefix.truncate(new_prefix_len);
111 }
112
113 /// Indicate that the next child node was ignored.
114 ///
115 /// This semantically combines start/finish operations.
116 pub fn ignore_leaf(&mut self, name: &str, is_last: bool) {
117 let has_columns = self.has_columns();
118
119 let buf = &mut self.write_buf;
120 buf.clear();
121
122 let branch = if !is_last { "├─ " } else { "╰─ " };
123 buf.extend([self.current_prefix.as_str(), branch, name]);
124
125 right_pad_buffer(buf, &mut self.max_name_span);
126
127 if has_columns {
128 TreeColumnData::from_first("(ignored)").write(buf, &mut self.column_widths);
129 } else {
130 buf.push_str("(ignored)");
131 }
132
133 println!("{buf}");
134 }
135
136 /// Enter a leaf node.
137 pub fn start_leaf(&mut self, name: &str, is_last: bool) {
138 let has_columns = self.has_columns();
139
140 let buf = &mut self.write_buf;
141 buf.clear();
142
143 let branch = if !is_last { "├─ " } else { "╰─ " };
144 buf.extend([self.current_prefix.as_str(), branch, name]);
145
146 // Right-pad buffer if this leaf will have info displayed.
147 if has_columns {
148 let max_span = self.max_name_span;
149 let buf_len = buf.chars().count();
150 let pad_len = TREE_COL_BUF + max_span.saturating_sub(buf_len);
151 buf.extend(repeat(' ').take(pad_len));
152
153 if buf_len > max_span {
154 self.max_name_span = buf_len;
155 }
156 }
157
158 print!("{buf}");
159 _ = std::io::stdout().flush();
160 }
161
162 /// Exit the current leaf node.
163 pub fn finish_empty_leaf(&mut self) {
164 println!();
165 }
166
167 /// Exit the current leaf node, emitting statistics.
168 pub fn finish_leaf(&mut self, is_last: bool, stats: &Stats, bytes_format: BytesFormat) {
169 let prep_buffer = |buf: &mut String, max_span: &mut usize| {
170 buf.clear();
171 buf.push_str(&self.current_prefix);
172
173 if !is_last {
174 buf.push('│');
175 }
176
177 right_pad_buffer(buf, max_span);
178 };
179
180 let buf = &mut self.write_buf;
181 buf.clear();
182
183 // Serialize max alloc counts and sizes early so we can resize columns
184 // early.
185 let serialized_max_alloc_counts = if stats.max_alloc.size.is_zero() {
186 None
187 } else {
188 Some(TreeColumn::ALL.map(|column| {
189 let Some(&max_alloc_count) = column.get_stat(&stats.max_alloc.count) else {
190 return String::new();
191 };
192
193 let prefix = if column.is_first() { " " } else { "" };
194 format!("{prefix}{}", util::fmt::format_f64(max_alloc_count, 4))
195 }))
196 };
197
198 let serialized_max_alloc_sizes = if stats.max_alloc.size.is_zero() {
199 None
200 } else {
201 Some(TreeColumn::ALL.map(|column| {
202 let Some(&max_alloc_size) = column.get_stat(&stats.max_alloc.size) else {
203 return String::new();
204 };
205
206 let prefix = if column.is_first() { " " } else { "" };
207 format!("{prefix}{}", util::fmt::format_bytes(max_alloc_size, 4, bytes_format))
208 }))
209 };
210
211 // Serialize alloc tallies early so we can resize columns early.
212 let serialized_alloc_tallies = AllocOp::ALL.map(|op| {
213 let tally = stats.alloc_tallies.get(op);
214
215 if tally.is_zero() {
216 return None;
217 }
218
219 let column_tallies = TreeColumn::ALL.map(|column| {
220 let prefix = if column.is_first() { " " } else { "" };
221
222 let tally = AllocTally {
223 count: column.get_stat(&tally.count).copied()?,
224 size: column.get_stat(&tally.size).copied()?,
225 };
226
227 Some((prefix, tally))
228 });
229
230 Some(AllocTally {
231 count: column_tallies.map(|tally| {
232 if let Some((prefix, tally)) = tally {
233 format!("{prefix}{}", util::fmt::format_f64(tally.count, 4))
234 } else {
235 String::new()
236 }
237 }),
238 size: column_tallies.map(|tally| {
239 if let Some((prefix, tally)) = tally {
240 format!("{prefix}{}", util::fmt::format_bytes(tally.size, 4, bytes_format))
241 } else {
242 String::new()
243 }
244 }),
245 })
246 });
247
248 // Serialize counter stats early so we can resize columns early.
249 let serialized_counters = KnownCounterKind::ALL.map(|counter_kind| {
250 let counter_stats = stats.get_counts(counter_kind);
251
252 TreeColumn::ALL
253 .map(|column| -> Option<String> {
254 let count = *column.get_stat(counter_stats?)?;
255 let time = *column.get_stat(&stats.time)?;
256
257 Some(
258 AnyCounter::known(counter_kind, count)
259 .display_throughput(time, bytes_format)
260 .to_string(),
261 )
262 })
263 .map(Option::unwrap_or_default)
264 });
265
266 // Set column widths based on serialized strings.
267 for column in TreeColumn::time_stats() {
268 let width = &mut self.column_widths[column as usize];
269
270 let mut update_width = |s: &str| {
271 *width = (*width).max(s.chars().count());
272 };
273
274 for counter in &serialized_counters {
275 update_width(&counter[column as usize]);
276 }
277
278 let serialized_max_alloc_counts = serialized_max_alloc_counts.iter().flatten();
279 let serialized_max_alloc_sizes = serialized_max_alloc_sizes.iter().flatten();
280 for s in serialized_max_alloc_counts.chain(serialized_max_alloc_sizes) {
281 update_width(s);
282 }
283
284 for s in serialized_alloc_tallies
285 .iter()
286 .flatten()
287 .flat_map(AllocTally::as_array)
288 .map(|values| &values[column as usize])
289 {
290 update_width(s);
291 }
292 }
293
294 // Write time stats with iter and sample counts.
295 TreeColumnData::from_fn(|column| -> String {
296 let stat: &dyn ToString = match column {
297 TreeColumn::Fastest => &stats.time.fastest,
298 TreeColumn::Slowest => &stats.time.slowest,
299 TreeColumn::Median => &stats.time.median,
300 TreeColumn::Mean => &stats.time.mean,
301 TreeColumn::Samples => &stats.sample_count,
302 TreeColumn::Iters => &stats.iter_count,
303 };
304 stat.to_string()
305 })
306 .as_ref::<str>()
307 .write(buf, &mut self.column_widths);
308
309 println!("{buf}");
310
311 // Write counter stats.
312 let counter_stats = serialized_counters.map(TreeColumnData);
313 for counter_kind in KnownCounterKind::ALL {
314 let counter_stats = counter_stats[counter_kind as usize].as_ref::<str>();
315
316 // Skip empty rows.
317 if counter_stats.0.iter().all(|s| s.is_empty()) {
318 continue;
319 }
320
321 prep_buffer(buf, &mut self.max_name_span);
322
323 counter_stats.write(buf, &mut self.column_widths);
324 println!("{buf}");
325 }
326
327 // Write max allocated bytes.
328 if serialized_max_alloc_counts.is_some() || serialized_max_alloc_sizes.is_some() {
329 prep_buffer(buf, &mut self.max_name_span);
330
331 TreeColumnData::from_first("max alloc:").write(buf, &mut self.column_widths);
332 println!("{buf}");
333
334 for serialized in
335 [serialized_max_alloc_counts.as_ref(), serialized_max_alloc_sizes.as_ref()]
336 .into_iter()
337 .flatten()
338 {
339 prep_buffer(buf, &mut self.max_name_span);
340
341 TreeColumnData::from_fn(|column| serialized[column as usize].as_str())
342 .write(buf, &mut self.column_widths);
343
344 println!("{buf}");
345 }
346 }
347
348 // Write allocation tallies.
349 for op in [AllocOp::Alloc, AllocOp::Dealloc, AllocOp::Grow, AllocOp::Shrink] {
350 let Some(tallies) = &serialized_alloc_tallies[op as usize] else {
351 continue;
352 };
353
354 prep_buffer(buf, &mut self.max_name_span);
355
356 TreeColumnData::from_first(op.prefix()).write(buf, &mut self.column_widths);
357 println!("{buf}");
358
359 for value in tallies.as_array() {
360 prep_buffer(buf, &mut self.max_name_span);
361
362 TreeColumnData::from_fn(|column| value[column as usize].as_str())
363 .write(buf, &mut self.column_widths);
364
365 println!("{buf}");
366 }
367 }
368 }
369
370 fn has_columns(&self) -> bool {
371 !self.column_widths.iter().all(|&w| w == 0)
372 }
373}
374
375/// Columns of the table next to the tree.
376#[derive(Clone, Copy, PartialEq, Eq)]
377pub(crate) enum TreeColumn {
378 Fastest,
379 Slowest,
380 Median,
381 Mean,
382 Samples,
383 Iters,
384}
385
386impl TreeColumn {
387 pub const COUNT: usize = 6;
388
389 pub const ALL: [Self; Self::COUNT] = {
390 use TreeColumn::*;
391 [Fastest, Slowest, Median, Mean, Samples, Iters]
392 };
393
394 #[inline]
395 pub fn time_stats() -> impl Iterator<Item = Self> {
396 use TreeColumn::*;
397 [Fastest, Slowest, Median, Mean].into_iter()
398 }
399
400 #[inline]
401 pub fn is_first(self) -> bool {
402 let [first, ..] = Self::ALL;
403 self == first
404 }
405
406 #[inline]
407 pub fn is_last(self) -> bool {
408 let [.., last] = Self::ALL;
409 self == last
410 }
411
412 fn name(self) -> &'static str {
413 match self {
414 Self::Fastest => "fastest",
415 Self::Slowest => "slowest",
416 Self::Median => "median",
417 Self::Mean => "mean",
418 Self::Samples => "samples",
419 Self::Iters => "iters",
420 }
421 }
422
423 #[inline]
424 pub fn is_time_stat(self) -> bool {
425 use TreeColumn::*;
426 matches!(self, Fastest | Slowest | Median | Mean)
427 }
428
429 #[inline]
430 fn get_stat<T>(self, stats: &StatsSet<T>) -> Option<&T> {
431 match self {
432 Self::Fastest => Some(&stats.fastest),
433 Self::Slowest => Some(&stats.slowest),
434 Self::Median => Some(&stats.median),
435 Self::Mean => Some(&stats.mean),
436 Self::Samples | Self::Iters => None,
437 }
438 }
439}
440
441#[derive(Default)]
442struct TreeColumnData<T>([T; TreeColumn::COUNT]);
443
444impl<T> TreeColumnData<T> {
445 #[inline]
446 fn from_first(value: T) -> Self
447 where
448 Self: Default,
449 {
450 let mut data: TreeColumnData = Self::default();
451 data.0[0] = value;
452 data
453 }
454
455 #[inline]
456 fn from_fn<F>(f: F) -> Self
457 where
458 F: FnMut(TreeColumn) -> T,
459 {
460 Self(TreeColumn::ALL.map(f))
461 }
462}
463
464impl TreeColumnData<&str> {
465 /// Writes the column data into the buffer.
466 fn write(&self, buf: &mut String, column_widths: &mut [usize; TreeColumn::COUNT]) {
467 for (column, value) in self.0.iter().enumerate() {
468 let is_first = column == 0;
469 let is_last = column == TreeColumn::COUNT - 1;
470
471 let value_width = value.chars().count();
472
473 // Write separator.
474 if !is_first {
475 let mut sep = " │ ";
476
477 // Prevent trailing spaces.
478 if is_last && value_width == 0 {
479 sep = &sep[..sep.len() - 1];
480 };
481
482 buf.push_str(sep);
483 }
484
485 buf.push_str(value);
486
487 // Right-pad remaining width or update column width to new maximum.
488 if !is_last {
489 if let Some(rem_width) = column_widths[column].checked_sub(value_width) {
490 buf.extend(repeat(' ').take(rem_width));
491 } else {
492 column_widths[column] = value_width;
493 }
494 }
495 }
496 }
497}
498
499impl<T> TreeColumnData<T> {
500 #[inline]
501 fn as_ref<U: ?Sized>(&self) -> TreeColumnData<&U>
502 where
503 T: AsRef<U>,
504 {
505 TreeColumnData::from_fn(|column: TreeColumn| self.0[column as usize].as_ref())
506 }
507}
508
509fn right_pad_buffer(buf: &mut String, max_span: &mut usize) {
510 let buf_len: usize = buf.chars().count();
511 let pad_len: usize = TREE_COL_BUF + max_span.saturating_sub(buf_len);
512 buf.extend(iter:repeat(elt:' ').take(pad_len));
513
514 if buf_len > *max_span {
515 *max_span = buf_len;
516 }
517}
518