Make stats code nicer.

Taking inspiration from `-Zmacro-stats`:
- Use "{prefix}" consistently.
- Use names for column widths.
- Write output in a single `eprint!` call, in an attempt to minimize
  interleaving of output from different rustc processes.
- Use `repeat` for the long `---` banners.

3 files changed, 62 insertions, 30 deletions

diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs
index 201b7e2b940..13c2a22172c 100644
--- a/compiler/rustc_interface/src/passes.rs
+++ b/compiler/rustc_interface/src/passes.rs
@@ -370,7 +370,7 @@ fn early_lint_checks(tcx: TyCtxt<'_>, (): ()) {
     let mut lint_buffer = resolver.lint_buffer.steal();
 
     if sess.opts.unstable_opts.input_stats {
-        input_stats::print_ast_stats(krate, "POST EXPANSION AST STATS", "ast-stats");
+        input_stats::print_ast_stats(krate);
     }
 
     // Needs to go *after* expansion to be able to check the results of macro expansion.
diff --git a/compiler/rustc_metadata/src/rmeta/encoder.rs b/compiler/rustc_metadata/src/rmeta/encoder.rs
index 00bd32eb0eb..53cf69b6ac2 100644
--- a/compiler/rustc_metadata/src/rmeta/encoder.rs
+++ b/compiler/rustc_metadata/src/rmeta/encoder.rs
@@ -762,6 +762,8 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
         assert_eq!(total_bytes, computed_total_bytes);
 
         if tcx.sess.opts.unstable_opts.meta_stats {
+            use std::fmt::Write;
+
             self.opaque.flush();
 
             // Rewind and re-read all the metadata to count the zero bytes we wrote.
@@ -781,27 +783,38 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             let prefix = "meta-stats";
             let perc = |bytes| (bytes * 100) as f64 / total_bytes as f64;
 
-            eprintln!("{prefix} METADATA STATS");
-            eprintln!("{} {:<23}{:>10}", prefix, "Section", "Size");
-            eprintln!("{prefix} ----------------------------------------------------------------");
+            let section_w = 23;
+            let size_w = 10;
+            let banner_w = 64;
+
+            // We write all the text into a string and print it with a single
+            // `eprint!`. This is an attempt to minimize interleaved text if multiple
+            // rustc processes are printing macro-stats at the same time (e.g. with
+            // `RUSTFLAGS='-Zmeta-stats' cargo build`). It still doesn't guarantee
+            // non-interleaving, though.
+            let mut s = String::new();
+            _ = writeln!(s, "{prefix} METADATA STATS");
+            _ = writeln!(s, "{prefix} {:<section_w$}{:>size_w$}", "Section", "Size");
+            _ = writeln!(s, "{prefix} {}", "-".repeat(banner_w));
             for (label, size) in stats {
-                eprintln!(
-                    "{} {:<23}{:>10} ({:4.1}%)",
-                    prefix,
+                _ = writeln!(
+                    s,
+                    "{prefix} {:<section_w$}{:>size_w$} ({:4.1}%)",
                     label,
                     usize_with_underscores(size),
                     perc(size)
                 );
             }
-            eprintln!("{prefix} ----------------------------------------------------------------");
-            eprintln!(
-                "{} {:<23}{:>10} (of which {:.1}% are zero bytes)",
-                prefix,
+            _ = writeln!(s, "{prefix} {}", "-".repeat(banner_w));
+            _ = writeln!(
+                s,
+                "{prefix} {:<section_w$}{:>size_w$} (of which {:.1}% are zero bytes)",
                 "Total",
                 usize_with_underscores(total_bytes),
                 perc(zero_bytes)
             );
-            eprintln!("{prefix}");
+            _ = writeln!(s, "{prefix}");
+            eprint!("{s}");
         }
 
         root
diff --git a/compiler/rustc_passes/src/input_stats.rs b/compiler/rustc_passes/src/input_stats.rs
index 46e6c0bf7da..8753e7e8e1f 100644
--- a/compiler/rustc_passes/src/input_stats.rs
+++ b/compiler/rustc_passes/src/input_stats.rs
@@ -68,13 +68,13 @@ pub fn print_hir_stats(tcx: TyCtxt<'_>) {
     collector.print("HIR STATS", "hir-stats");
 }
 
-pub fn print_ast_stats(krate: &ast::Crate, title: &str, prefix: &str) {
+pub fn print_ast_stats(krate: &ast::Crate) {
     use rustc_ast::visit::Visitor;
 
     let mut collector =
         StatCollector { tcx: None, nodes: FxHashMap::default(), seen: FxHashSet::default() };
     collector.visit_crate(krate);
-    collector.print(title, prefix);
+    collector.print("POST EXPANSION AST STATS", "ast-stats");
 }
 
 impl<'k> StatCollector<'k> {
@@ -117,28 +117,45 @@ impl<'k> StatCollector<'k> {
     }
 
     fn print(&self, title: &str, prefix: &str) {
+        use std::fmt::Write;
+
         // We will soon sort, so the initial order does not matter.
         #[allow(rustc::potential_query_instability)]
         let mut nodes: Vec<_> = self.nodes.iter().collect();
         nodes.sort_by_cached_key(|(label, node)| (node.stats.accum_size(), label.to_owned()));
 
+        let name_w = 18;
+        let acc_size1_w = 10;
+        let acc_size2_w = 8; // " (NN.N%)"
+        let acc_size_w = acc_size1_w + acc_size2_w;
+        let count_w = 14;
+        let item_size_w = 14;
+        let banner_w = name_w + acc_size_w + count_w + item_size_w;
+
         let total_size = nodes.iter().map(|(_, node)| node.stats.accum_size()).sum();
         let total_count = nodes.iter().map(|(_, node)| node.stats.count).sum();
 
-        eprintln!("{prefix} {title}");
-        eprintln!(
-            "{} {:<18}{:>18}{:>14}{:>14}",
-            prefix, "Name", "Accumulated Size", "Count", "Item Size"
+        // We write all the text into a string and print it with a single
+        // `eprint!`. This is an attempt to minimize interleaved text if multiple
+        // rustc processes are printing macro-stats at the same time (e.g. with
+        // `RUSTFLAGS='-Zinput-stats' cargo build`). It still doesn't guarantee
+        // non-interleaving, though.
+        let mut s = String::new();
+        _ = writeln!(s, "{prefix} {title}");
+        _ = writeln!(
+            s,
+            "{prefix} {:<name_w$}{:>acc_size_w$}{:>count_w$}{:>item_size_w$}",
+            "Name", "Accumulated Size", "Count", "Item Size"
         );
-        eprintln!("{prefix} ----------------------------------------------------------------");
+        _ = writeln!(s, "{prefix} {}", "-".repeat(banner_w));
 
         let percent = |m, n| (m * 100) as f64 / n as f64;
 
         for (label, node) in nodes {
             let size = node.stats.accum_size();
-            eprintln!(
-                "{} {:<18}{:>10} ({:4.1}%){:>14}{:>14}",
-                prefix,
+            _ = writeln!(
+                s,
+                "{prefix} {:<name_w$}{:>acc_size1_w$} ({:4.1}%){:>count_w$}{:>item_size_w$}",
                 label,
                 usize_with_underscores(size),
                 percent(size, total_size),
@@ -155,9 +172,9 @@ impl<'k> StatCollector<'k> {
 
                 for (label, subnode) in subnodes {
                     let size = subnode.accum_size();
-                    eprintln!(
-                        "{} - {:<18}{:>10} ({:4.1}%){:>14}",
-                        prefix,
+                    _ = writeln!(
+                        s,
+                        "{prefix} - {:<name_w$}{:>acc_size1_w$} ({:4.1}%){:>count_w$}",
                         label,
                         usize_with_underscores(size),
                         percent(size, total_size),
@@ -166,15 +183,17 @@ impl<'k> StatCollector<'k> {
                 }
             }
         }
-        eprintln!("{prefix} ----------------------------------------------------------------");
-        eprintln!(
-            "{} {:<18}{:>10}        {:>14}",
-            prefix,
+        _ = writeln!(s, "{prefix} {}", "-".repeat(banner_w));
+        _ = writeln!(
+            s,
+            "{prefix} {:<name_w$}{:>acc_size1_w$}{:>acc_size2_w$}{:>count_w$}",
             "Total",
             usize_with_underscores(total_size),
+            "",
             usize_with_underscores(total_count),
         );
-        eprintln!("{prefix}");
+        _ = writeln!(s, "{prefix}");
+        eprint!("{s}");
     }
 }