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

8 files changed, 151 insertions, 5 deletions

diff --git a/compiler/rustc_driver_impl/src/lib.rs b/compiler/rustc_driver_impl/src/lib.rs
index 5b75205442b..f8dda23ef9c 100644
--- a/compiler/rustc_driver_impl/src/lib.rs
+++ b/compiler/rustc_driver_impl/src/lib.rs
@@ -430,6 +430,13 @@ fn run_compiler(
                 sess.code_stats.print_type_sizes();
             }
 
+            if sess.opts.unstable_opts.print_vtable_sizes {
+                let crate_name =
+                    compiler.session().opts.crate_name.as_deref().unwrap_or("<UNKNOWN_CRATE>");
+
+                sess.code_stats.print_vtable_sizes(crate_name);
+            }
+
             let linker = queries.linker()?;
             Ok(Some(linker))
         })?;
diff --git a/compiler/rustc_interface/src/interface.rs b/compiler/rustc_interface/src/interface.rs
index 2edc72ba72e..5e3ea71f0e7 100644
--- a/compiler/rustc_interface/src/interface.rs
+++ b/compiler/rustc_interface/src/interface.rs
@@ -333,6 +333,7 @@ pub fn run_compiler<R: Send>(config: Config, f: impl FnOnce(&Compiler) -> R + Se
                 };
 
                 let prof = compiler.sess.prof.clone();
+
                 prof.generic_activity("drop_compiler").run(move || drop(compiler));
                 r
             })
diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs
index 83a74742f5b..be6bd913790 100644
--- a/compiler/rustc_interface/src/passes.rs
+++ b/compiler/rustc_interface/src/passes.rs
@@ -24,6 +24,7 @@ use rustc_parse::{parse_crate_from_file, parse_crate_from_source_str, validate_a
 use rustc_passes::{self, hir_stats, layout_test};
 use rustc_plugin_impl as plugin;
 use rustc_resolve::Resolver;
+use rustc_session::code_stats::VTableSizeInfo;
 use rustc_session::config::{CrateType, Input, OutFileName, OutputFilenames, OutputType};
 use rustc_session::cstore::{MetadataLoader, Untracked};
 use rustc_session::output::filename_for_input;
@@ -866,6 +867,99 @@ fn analysis(tcx: TyCtxt<'_>, (): ()) -> Result<()> {
         sess.time("check_lint_expectations", || tcx.check_expectations(None));
     });
 
+    if sess.opts.unstable_opts.print_vtable_sizes {
+        let traits = tcx.traits(LOCAL_CRATE);
+
+        for &tr in traits {
+            if !tcx.check_is_object_safe(tr) {
+                continue;
+            }
+
+            let name = ty::print::with_no_trimmed_paths!(tcx.def_path_str(tr));
+
+            let mut first_dsa = true;
+
+            // Number of vtable entries, if we didn't have upcasting
+            let mut unupcasted_cost = 0;
+            // Number of vtable entries needed solely for upcasting
+            let mut upcast_cost = 0;
+
+            let trait_ref = ty::Binder::dummy(ty::TraitRef::identity(tcx, tr));
+
+            // A slightly edited version of the code in `rustc_trait_selection::traits::vtable::vtable_entries`,
+            // that works without self type and just counts number of entries.
+            //
+            // Note that this is technically wrong, for traits which have associated types in supertraits:
+            //
+            //   trait A: AsRef<Self::T> + AsRef<()> { type T; }
+            //
+            // Without self type we can't normalize `Self::T`, so we can't know if `AsRef<Self::T>` and
+            // `AsRef<()>` are the same trait, thus we assume that those are different, and potentially
+            // over-estimate how many vtable entries there are.
+            //
+            // Similarly this is wrong for traits that have methods with possibly-impossible bounds.
+            // For example:
+            //
+            //   trait B<T> { fn f(&self) where T: Copy; }
+            //
+            // Here `dyn B<u8>` will have 4 entries, while `dyn B<String>` will only have 3.
+            // However, since we don't know `T`, we can't know if `T: Copy` holds or not,
+            // thus we lean on the bigger side and say it has 4 entries.
+            traits::vtable::prepare_vtable_segments(tcx, trait_ref, |segment| {
+                match segment {
+                    traits::vtable::VtblSegment::MetadataDSA => {
+                        // If this is the first dsa, it would be included either way,
+                        // otherwise it's needed for upcasting
+                        if std::mem::take(&mut first_dsa) {
+                            unupcasted_cost += 3;
+                        } else {
+                            upcast_cost += 3;
+                        }
+                    }
+
+                    traits::vtable::VtblSegment::TraitOwnEntries { trait_ref, emit_vptr } => {
+                        let existential_trait_ref = trait_ref.map_bound(|trait_ref| {
+                            ty::ExistentialTraitRef::erase_self_ty(tcx, trait_ref)
+                        });
+
+                        // Lookup the shape of vtable for the trait.
+                        let own_existential_entries =
+                            tcx.own_existential_vtable_entries(existential_trait_ref.def_id());
+
+                        let own_entries = own_existential_entries.iter().copied().map(|_def_id| {
+                            // The original code here ignores the method if its predicates are impossible.
+                            // We can't really do that as, for example, all not trivial bounds on generic
+                            // parameters are impossible (since we don't know the parameters...),
+                            // see the comment above.
+
+                            1
+                        });
+
+                        unupcasted_cost += own_entries.sum::<usize>();
+
+                        if emit_vptr {
+                            upcast_cost += 1;
+                        }
+                    }
+                }
+
+                std::ops::ControlFlow::Continue::<std::convert::Infallible>(())
+            });
+
+            sess.code_stats.record_vtable_size(
+                tr,
+                &name,
+                VTableSizeInfo {
+                    trait_name: name.clone(),
+                    size_words_without_upcasting: unupcasted_cost,
+                    size_words_with_upcasting: unupcasted_cost + upcast_cost,
+                    difference_words: upcast_cost,
+                    difference_percent: upcast_cost as f64 / unupcasted_cost as f64 * 100.,
+                },
+            )
+        }
+    }
+
     Ok(())
 }
 
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,
diff --git a/compiler/rustc_trait_selection/src/traits/mod.rs b/compiler/rustc_trait_selection/src/traits/mod.rs
index c2f94cb6385..ab58684ffd9 100644
--- a/compiler/rustc_trait_selection/src/traits/mod.rs
+++ b/compiler/rustc_trait_selection/src/traits/mod.rs
@@ -21,7 +21,7 @@ mod structural_match;
 mod structural_normalize;
 #[cfg_attr(not(bootstrap), allow(hidden_glob_reexports))]
 mod util;
-mod vtable;
+pub mod vtable;
 pub mod wf;
 
 use crate::infer::outlives::env::OutlivesEnvironment;
diff --git a/compiler/rustc_trait_selection/src/traits/vtable.rs b/compiler/rustc_trait_selection/src/traits/vtable.rs
index cc674ceee3d..96b8e0b82b6 100644
--- a/compiler/rustc_trait_selection/src/traits/vtable.rs
+++ b/compiler/rustc_trait_selection/src/traits/vtable.rs
@@ -15,13 +15,13 @@ use std::fmt::Debug;
 use std::ops::ControlFlow;
 
 #[derive(Clone, Debug)]
-pub(super) enum VtblSegment<'tcx> {
+pub enum VtblSegment<'tcx> {
     MetadataDSA,
     TraitOwnEntries { trait_ref: ty::PolyTraitRef<'tcx>, emit_vptr: bool },
 }
 
 /// Prepare the segments for a vtable
-pub(super) fn prepare_vtable_segments<'tcx, T>(
+pub fn prepare_vtable_segments<'tcx, T>(
     tcx: TyCtxt<'tcx>,
     trait_ref: ty::PolyTraitRef<'tcx>,
     mut segment_visitor: impl FnMut(VtblSegment<'tcx>) -> ControlFlow<T>,