From 12eaa3ab8465431d18696489e29be95fedf4da6d Mon Sep 17 00:00:00 2001
From: Sam Estep <sam@samestep.com>
Date: Wed, 13 Nov 2024 23:19:47 -0500
Subject: [PATCH] Print total node count in `-Z hir-stats`

---
 compiler/rustc_passes/src/hir_stats.rs | 9 ++++++++-
 tests/ui/stats/hir-stats.rs            | 2 +-
 tests/ui/stats/hir-stats.stderr        | 6 +++---
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/compiler/rustc_passes/src/hir_stats.rs b/compiler/rustc_passes/src/hir_stats.rs
index 55fdbac8ad6..db34189be2a 100644
--- a/compiler/rustc_passes/src/hir_stats.rs
+++ b/compiler/rustc_passes/src/hir_stats.rs
@@ -126,6 +126,7 @@ impl<'k> StatCollector<'k> {
         });
 
         let total_size = nodes.iter().map(|(_, node)| node.stats.count * node.stats.size).sum();
+        let total_count = nodes.iter().map(|(_, node)| node.stats.count).sum();
 
         eprintln!("{prefix} {title}");
         eprintln!(
@@ -167,7 +168,13 @@ impl<'k> StatCollector<'k> {
             }
         }
         eprintln!("{prefix} ----------------------------------------------------------------");
-        eprintln!("{} {:<18}{:>10}", prefix, "Total", to_readable_str(total_size));
+        eprintln!(
+            "{} {:<18}{:>10}        {:>14}",
+            prefix,
+            "Total",
+            to_readable_str(total_size),
+            to_readable_str(total_count),
+        );
         eprintln!("{prefix}");
     }
 }
diff --git a/tests/ui/stats/hir-stats.rs b/tests/ui/stats/hir-stats.rs
index 7c5da8cf554..dc4a90c0059 100644
--- a/tests/ui/stats/hir-stats.rs
+++ b/tests/ui/stats/hir-stats.rs
@@ -1,6 +1,6 @@
 //@ check-pass
 //@ compile-flags: -Zhir-stats
-//@ only-x86_64
+//@ only-64bit
 // layout randomization affects the hir stat output
 //@ needs-deterministic-layouts
 
diff --git a/tests/ui/stats/hir-stats.stderr b/tests/ui/stats/hir-stats.stderr
index bd0c8cbc3b1..2cc2c8019fa 100644
--- a/tests/ui/stats/hir-stats.stderr
+++ b/tests/ui/stats/hir-stats.stderr
@@ -53,7 +53,7 @@ ast-stats-1 - Enum                     136 ( 2.0%)             1
 ast-stats-1 - Fn                       272 ( 4.1%)             2
 ast-stats-1 - Use                      408 ( 6.1%)             3
 ast-stats-1 ----------------------------------------------------------------
-ast-stats-1 Total                  6_640
+ast-stats-1 Total                  6_640                   116
 ast-stats-1
 ast-stats-2 POST EXPANSION AST STATS
 ast-stats-2 Name                Accumulated Size         Count     Item Size
@@ -113,7 +113,7 @@ ast-stats-2 - ForeignMod               136 ( 1.9%)             1
 ast-stats-2 - Fn                       272 ( 3.7%)             2
 ast-stats-2 - Use                      544 ( 7.5%)             4
 ast-stats-2 ----------------------------------------------------------------
-ast-stats-2 Total                  7_288
+ast-stats-2 Total                  7_288                   127
 ast-stats-2
 hir-stats HIR STATS
 hir-stats Name                Accumulated Size         Count     Item Size
@@ -174,5 +174,5 @@ hir-stats - Use                      352 ( 3.9%)             4
 hir-stats Path                   1_240 (13.7%)            31            40
 hir-stats PathSegment            1_920 (21.3%)            40            48
 hir-stats ----------------------------------------------------------------
-hir-stats Total                  9_024
+hir-stats Total                  9_024                   180
 hir-stats