Auto merge of #88552 - nbdd0121:vtable, r=nagisa

Stop allocating vtable entries for non-object-safe methods

Current a vtable entry is allocated for all associated fns, even if the method is not object-safe: https://godbolt.org/z/h7vx6f35T

As a result, each vtable for `Iterator`' currently consumes 74 `usize`s. This PR stops allocating vtable entries for those methods, reducing vtable size of each `Iterator` vtable to 7 `usize`s.

Note that this PR introduces will cause more invocations of `is_vtable_safe_method`. So a perf run might be needed. If result isn't favorable then we might need to query-ify `is_vtable_safe_method`.

8 files changed, 109 insertions, 42 deletions

diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index dada94edc95..c93996162e3 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -996,6 +996,12 @@ rustc_queries! {
         desc { |tcx| "checking if item has mir available: `{}`", tcx.def_path_str(key) }
     }
 
+    query own_existential_vtable_entries(
+        key: ty::PolyExistentialTraitRef<'tcx>
+    ) -> &'tcx [DefId] {
+        desc { |tcx| "finding all existential vtable entries for trait {}", tcx.def_path_str(key.def_id()) }
+    }
+
     query vtable_entries(key: ty::PolyTraitRef<'tcx>)
                         -> &'tcx [ty::VtblEntry<'tcx>] {
         desc { |tcx| "finding all vtable entries for trait {}", tcx.def_path_str(key.def_id()) }
diff --git a/compiler/rustc_query_impl/src/keys.rs b/compiler/rustc_query_impl/src/keys.rs
index c973eae6b06..42e8b4023cf 100644
--- a/compiler/rustc_query_impl/src/keys.rs
+++ b/compiler/rustc_query_impl/src/keys.rs
@@ -294,6 +294,16 @@ impl<'tcx> Key for ty::PolyTraitRef<'tcx> {
     }
 }
 
+impl<'tcx> Key for ty::PolyExistentialTraitRef<'tcx> {
+    #[inline(always)]
+    fn query_crate_is_local(&self) -> bool {
+        self.def_id().krate == LOCAL_CRATE
+    }
+    fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
+        tcx.def_span(self.def_id())
+    }
+}
+
 impl<'tcx> Key for (ty::PolyTraitRef<'tcx>, ty::PolyTraitRef<'tcx>) {
     #[inline(always)]
     fn query_crate_is_local(&self) -> bool {
diff --git a/compiler/rustc_trait_selection/src/traits/mod.rs b/compiler/rustc_trait_selection/src/traits/mod.rs
index 17a4184c3c9..44c67524383 100644
--- a/compiler/rustc_trait_selection/src/traits/mod.rs
+++ b/compiler/rustc_trait_selection/src/traits/mod.rs
@@ -625,6 +625,31 @@ fn dump_vtable_entries<'tcx>(
     tcx.sess.struct_span_err(sp, &msg).emit();
 }
 
+fn own_existential_vtable_entries<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    trait_ref: ty::PolyExistentialTraitRef<'tcx>,
+) -> &'tcx [DefId] {
+    let trait_methods = tcx
+        .associated_items(trait_ref.def_id())
+        .in_definition_order()
+        .filter(|item| item.kind == ty::AssocKind::Fn);
+    // Now list each method's DefId (for within its trait).
+    let own_entries = trait_methods.filter_map(move |trait_method| {
+        debug!("own_existential_vtable_entry: trait_method={:?}", trait_method);
+        let def_id = trait_method.def_id;
+
+        // Some methods cannot be called on an object; skip those.
+        if !is_vtable_safe_method(tcx, trait_ref.def_id(), &trait_method) {
+            debug!("own_existential_vtable_entry: not vtable safe");
+            return None;
+        }
+
+        Some(def_id)
+    });
+
+    tcx.arena.alloc_from_iter(own_entries.into_iter())
+}
+
 /// Given a trait `trait_ref`, iterates the vtable entries
 /// that come from `trait_ref`, including its supertraits.
 fn vtable_entries<'tcx>(
@@ -641,21 +666,15 @@ fn vtable_entries<'tcx>(
                 entries.extend(COMMON_VTABLE_ENTRIES);
             }
             VtblSegment::TraitOwnEntries { trait_ref, emit_vptr } => {
-                let trait_methods = tcx
-                    .associated_items(trait_ref.def_id())
-                    .in_definition_order()
-                    .filter(|item| item.kind == ty::AssocKind::Fn);
-                // Now list each method's DefId and InternalSubsts (for within its trait).
-                // If the method can never be called from this object, produce `Vacant`.
-                let own_entries = trait_methods.map(move |trait_method| {
-                    debug!("vtable_entries: trait_method={:?}", trait_method);
-                    let def_id = trait_method.def_id;
-
-                    // Some methods cannot be called on an object; skip those.
-                    if !is_vtable_safe_method(tcx, trait_ref.def_id(), &trait_method) {
-                        debug!("vtable_entries: not vtable safe");
-                        return VtblEntry::Vacant;
-                    }
+                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);
+
+                let own_entries = own_existential_entries.iter().copied().map(|def_id| {
+                    debug!("vtable_entries: trait_method={:?}", def_id);
 
                     // The method may have some early-bound lifetimes; add regions for those.
                     let substs = trait_ref.map_bound(|trait_ref| {
@@ -804,6 +823,7 @@ pub fn provide(providers: &mut ty::query::Providers) {
         specialization_graph_of: specialize::specialization_graph_provider,
         specializes: specialize::specializes,
         codegen_fulfill_obligation: codegen::codegen_fulfill_obligation,
+        own_existential_vtable_entries,
         vtable_entries,
         vtable_trait_upcasting_coercion_new_vptr_slot,
         subst_and_check_impossible_predicates,
diff --git a/compiler/rustc_trait_selection/src/traits/util.rs b/compiler/rustc_trait_selection/src/traits/util.rs
index fd94f9f7998..b108d85bb20 100644
--- a/compiler/rustc_trait_selection/src/traits/util.rs
+++ b/compiler/rustc_trait_selection/src/traits/util.rs
@@ -285,15 +285,10 @@ pub fn upcast_choices(
 /// that come from `trait_ref`, excluding its supertraits. Used in
 /// computing the vtable base for an upcast trait of a trait object.
 pub fn count_own_vtable_entries(tcx: TyCtxt<'tcx>, trait_ref: ty::PolyTraitRef<'tcx>) -> usize {
-    let mut entries = 0;
-    // Count number of methods and add them to the total offset.
-    // Skip over associated types and constants.
-    for trait_item in tcx.associated_items(trait_ref.def_id()).in_definition_order() {
-        if trait_item.kind == ty::AssocKind::Fn {
-            entries += 1;
-        }
-    }
-    entries
+    let existential_trait_ref =
+        trait_ref.map_bound(|trait_ref| ty::ExistentialTraitRef::erase_self_ty(tcx, trait_ref));
+    let existential_trait_ref = tcx.erase_regions(existential_trait_ref);
+    tcx.own_existential_vtable_entries(existential_trait_ref).len()
 }
 
 /// Given an upcast trait object described by `object`, returns the
@@ -304,22 +299,21 @@ pub fn get_vtable_index_of_object_method<N>(
     object: &super::ImplSourceObjectData<'tcx, N>,
     method_def_id: DefId,
 ) -> usize {
+    let existential_trait_ref = object
+        .upcast_trait_ref
+        .map_bound(|trait_ref| ty::ExistentialTraitRef::erase_self_ty(tcx, trait_ref));
+    let existential_trait_ref = tcx.erase_regions(existential_trait_ref);
     // Count number of methods preceding the one we are selecting and
     // add them to the total offset.
-    // Skip over associated types and constants, as those aren't stored in the vtable.
-    let mut entries = object.vtable_base;
-    for trait_item in tcx.associated_items(object.upcast_trait_ref.def_id()).in_definition_order() {
-        if trait_item.def_id == method_def_id {
-            // The item with the ID we were given really ought to be a method.
-            assert_eq!(trait_item.kind, ty::AssocKind::Fn);
-            return entries;
-        }
-        if trait_item.kind == ty::AssocKind::Fn {
-            entries += 1;
-        }
-    }
-
-    bug!("get_vtable_index_of_object_method: {:?} was not found", method_def_id);
+    let index = tcx
+        .own_existential_vtable_entries(existential_trait_ref)
+        .iter()
+        .copied()
+        .position(|def_id| def_id == method_def_id)
+        .unwrap_or_else(|| {
+            bug!("get_vtable_index_of_object_method: {:?} was not found", method_def_id);
+        });
+    object.vtable_base + index
 }
 
 pub fn closure_trait_ref_and_return_type(
diff --git a/src/test/ui/traits/vtable/vtable-non-object-safe.rs b/src/test/ui/traits/vtable/vtable-non-object-safe.rs
new file mode 100644
index 00000000000..45b6a8a98a7
--- /dev/null
+++ b/src/test/ui/traits/vtable/vtable-non-object-safe.rs
@@ -0,0 +1,18 @@
+// build-fail
+#![feature(rustc_attrs)]
+
+// Ensure that non-object-safe methods in Iterator does not generate
+// vtable entries.
+
+#[rustc_dump_vtable]
+trait A: Iterator {}
+//~^ error Vtable
+
+impl<T> A for T where T: Iterator {}
+
+fn foo(_a: &mut dyn A<Item=u8>) {
+}
+
+fn main() {
+    foo(&mut vec![0, 1, 2, 3].into_iter());
+}
diff --git a/src/test/ui/traits/vtable/vtable-non-object-safe.stderr b/src/test/ui/traits/vtable/vtable-non-object-safe.stderr
new file mode 100644
index 00000000000..f3175b805d1
--- /dev/null
+++ b/src/test/ui/traits/vtable/vtable-non-object-safe.stderr
@@ -0,0 +1,16 @@
+error: Vtable entries for `<std::vec::IntoIter<u8> as A>`: [
+    MetadataDropInPlace,
+    MetadataSize,
+    MetadataAlign,
+    Method(<std::vec::IntoIter<u8> as Iterator>::next),
+    Method(<std::vec::IntoIter<u8> as Iterator>::size_hint),
+    Method(<std::vec::IntoIter<u8> as Iterator>::advance_by),
+    Method(<std::vec::IntoIter<u8> as Iterator>::nth),
+]
+  --> $DIR/vtable-non-object-safe.rs:8:1
+   |
+LL | trait A: Iterator {}
+   | ^^^^^^^^^^^^^^^^^^^^
+
+error: aborting due to previous error
+
diff --git a/src/test/ui/traits/vtable/vtable-vacant.rs b/src/test/ui/traits/vtable/vtable-vacant.rs
index ebea94171f2..429ce523799 100644
--- a/src/test/ui/traits/vtable/vtable-vacant.rs
+++ b/src/test/ui/traits/vtable/vtable-vacant.rs
@@ -1,22 +1,25 @@
 // build-fail
 #![feature(rustc_attrs)]
+#![feature(negative_impls)]
+#![allow(where_clauses_object_safety)]
 
 // B --> A
 
 #[rustc_dump_vtable]
 trait A {
     fn foo_a1(&self) {}
-    fn foo_a2(&self) where Self: Sized {}
+    fn foo_a2(&self) where Self: Send {}
 }
 
 #[rustc_dump_vtable]
 trait B: A {
     //~^ error Vtable
     fn foo_b1(&self) {}
-    fn foo_b2() where Self: Sized {}
+    fn foo_b2(&self) where Self: Send {}
 }
 
 struct S;
+impl !Send for S {}
 
 impl A for S {}
 impl B for S {}
diff --git a/src/test/ui/traits/vtable/vtable-vacant.stderr b/src/test/ui/traits/vtable/vtable-vacant.stderr
index 768cca52689..f5cd36264fc 100644
--- a/src/test/ui/traits/vtable/vtable-vacant.stderr
+++ b/src/test/ui/traits/vtable/vtable-vacant.stderr
@@ -7,12 +7,12 @@ error: Vtable entries for `<S as B>`: [
     Method(<S as B>::foo_b1),
     Vacant,
 ]
-  --> $DIR/vtable-vacant.rs:13:1
+  --> $DIR/vtable-vacant.rs:15:1
    |
 LL | / trait B: A {
 LL | |
 LL | |     fn foo_b1(&self) {}
-LL | |     fn foo_b2() where Self: Sized {}
+LL | |     fn foo_b2(&self) where Self: Send {}
 LL | | }
    | |_^