Collect VTable stats & add `-Zprint-vtable-sizes`

3 files changed, 46 insertions, 2 deletions

diff --git a/compiler/rustc_session/src/code_stats.rs b/compiler/rustc_session/src/code_stats.rs
index 0dfee92f404..f76263de13f 100644
--- a/compiler/rustc_session/src/code_stats.rs
+++ b/compiler/rustc_session/src/code_stats.rs
@@ -1,5 +1,6 @@
-use rustc_data_structures::fx::FxHashSet;
+use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::sync::Lock;
+use rustc_span::def_id::DefId;
 use rustc_span::Symbol;
 use rustc_target::abi::{Align, Size};
 use std::cmp;
@@ -65,9 +66,18 @@ pub struct TypeSizeInfo {
     pub variants: Vec<VariantInfo>,
 }
 
+pub struct VTableSizeInfo {
+    pub trait_name: String,
+    pub size_words_without_upcasting: usize,
+    pub size_words_with_upcasting: usize,
+    pub difference_words: usize,
+    pub difference_percent: f64,
+}
+
 #[derive(Default)]
 pub struct CodeStats {
     type_sizes: Lock<FxHashSet<TypeSizeInfo>>,
+    vtable_sizes: Lock<FxHashMap<DefId, VTableSizeInfo>>,
 }
 
 impl CodeStats {
@@ -101,6 +111,14 @@ impl CodeStats {
         self.type_sizes.borrow_mut().insert(info);
     }
 
+    pub fn record_vtable_size(&self, trait_did: DefId, trait_name: &str, info: VTableSizeInfo) {
+        let prev = self.vtable_sizes.lock().insert(trait_did, info);
+        assert!(
+            prev.is_none(),
+            "size of vtable for `{trait_name}` ({trait_did:?}) is already recorded"
+        );
+    }
+
     pub fn print_type_sizes(&self) {
         let type_sizes = self.type_sizes.borrow();
         let mut sorted: Vec<_> = type_sizes.iter().collect();
@@ -196,4 +214,28 @@ impl CodeStats {
             }
         }
     }
+
+    pub fn print_vtable_sizes(&self, crate_name: &str) {
+        let mut rr = std::mem::take(&mut *self.vtable_sizes.lock()).into_iter().collect::<Vec<_>>();
+
+        rr.sort_by(|(_, stats_a), (_, stats_b)| {
+            stats_b.difference_percent.total_cmp(&stats_a.difference_percent)
+        });
+
+        for (
+            _,
+            VTableSizeInfo {
+                trait_name,
+                size_words_without_upcasting,
+                size_words_with_upcasting,
+                difference_words,
+                difference_percent,
+            },
+        ) in rr
+        {
+            println!(
+                r#"print-vtable-sizes {{ "crate_name": "{crate_name}", "trait_name": "{trait_name}", "size_unupcastable_words": "{size_words_without_upcasting}", "size_upcastable_words": "{size_words_with_upcasting}", diff: "{difference_words}", diff_p: "{difference_percent}" }}"#
+            );
+        }
+    }
 }
diff --git a/compiler/rustc_session/src/lib.rs b/compiler/rustc_session/src/lib.rs
index 590a68c6600..d57aa820fcb 100644
--- a/compiler/rustc_session/src/lib.rs
+++ b/compiler/rustc_session/src/lib.rs
@@ -27,7 +27,7 @@ pub use lint::{declare_lint, declare_lint_pass, declare_tool_lint, impl_lint_pas
 pub use rustc_lint_defs as lint;
 pub mod parse;
 
-mod code_stats;
+pub mod code_stats;
 #[macro_use]
 pub mod config;
 pub mod cstore;
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
index 201066e3950..b626c721481 100644
--- a/compiler/rustc_session/src/options.rs
+++ b/compiler/rustc_session/src/options.rs
@@ -1632,6 +1632,8 @@ options! {
         "print the result of the monomorphization collection pass"),
     print_type_sizes: bool = (false, parse_bool, [UNTRACKED],
         "print layout information for each type encountered (default: no)"),
+    print_vtable_sizes: bool = (false, parse_bool, [UNTRACKED],
+        "print size comparison between old and new vtable layouts (default: no)"),
     proc_macro_backtrace: bool = (false, parse_bool, [UNTRACKED],
          "show backtraces for panics during proc-macro execution (default: no)"),
     proc_macro_execution_strategy: ProcMacroExecutionStrategy = (ProcMacroExecutionStrategy::SameThread,