Add a !type annotation to a vtable for every supertrait that is a prefix of the vtable

compiler/rustc_codegen_llvm/src/context.rs

@@ -941,11 +941,10 @@ impl<'ll, 'tcx> MiscCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {
 
     fn apply_vcall_visibility_metadata(
         &self,
-        ty: Ty<'tcx>,
-        poly_trait_ref: Option<ty::ExistentialTraitRef<'tcx>>,
+        trait_hierarchy: &[ty::ExistentialTraitRef<'tcx>],
         vtable: &'ll Value,
     ) {
-        apply_vcall_visibility_metadata(self, ty, poly_trait_ref, vtable);
+        apply_vcall_visibility_metadata(self, trait_hierarchy, vtable);
     }
 
     fn get_fn(&self, instance: Instance<'tcx>) -> &'ll Value {

compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs

@@ -1733,10 +1733,56 @@ fn find_vtable_behind_cast<'ll>(vtable: &'ll Value) -> &'ll Value {
     vtable
 }
 
+/// From where can virtual calls using this vtable be made?
+///
+/// See also <https://llvm.org/docs/TypeMetadata.html#vcall-visibility-metadata>.
+#[derive(PartialEq, Eq, PartialOrd, Ord)]
+enum VCallVisibility {
+    /// Virtual function calls using this vtable could be made from external code.
+    Public = 0,
+    /// All virtual function calls which might use this vtable are in the current LTO unit,
+    /// meaning they will be in the current module once LTO linking has been performed.
+    LinkageUnit = 1,
+    /// All virtual function calls which might use this vtable are in the current module.
+    TranslationUnit = 2,
+}
+
+impl VCallVisibility {
+    fn for_def_id<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>, def_id: DefId) -> Self {
+        let lto = cx.sess().lto();
+        let trait_vis = cx.tcx.visibility(def_id);
+        let cgus = cx.sess().codegen_units().as_usize();
+        let single_cgu = cgus == 1;
+
+        // Because LLVM requires full LTO for the virtual function elimination optimization to
+        // apply, only the `Lto::Fat` cases are relevant currently.
+        match (lto, trait_vis, single_cgu) {
+            // If there is not LTO and the visibility in public, we have to assume that the vtable
+            // can be seen from anywhere. With multiple CGUs, the vtable is quasi-public.
+            (Lto::No | Lto::ThinLocal, Visibility::Public, _)
+            | (Lto::No, Visibility::Restricted(_), false) => VCallVisibility::Public,
+            // With LTO and a quasi-public visibility, the usages of the functions of the vtable are
+            // all known by the `LinkageUnit`.
+            // FIXME: LLVM only supports this optimization for `Lto::Fat` currently. Once it also
+            // supports `Lto::Thin` the `VCallVisibility` may have to be adjusted for those.
+            (Lto::Fat | Lto::Thin, Visibility::Public, _)
+            | (Lto::ThinLocal | Lto::Thin | Lto::Fat, Visibility::Restricted(_), false) => {
+                VCallVisibility::LinkageUnit
+            }
+            // If there is only one CGU, private vtables can only be seen by that CGU/translation
+            // unit and therefore we know of all usages of functions in the vtable.
+            (_, Visibility::Restricted(_), true) => VCallVisibility::TranslationUnit,
+        }
+    }
+}
+
+/// Apply the virtual call visibility metadata to the vtable.
+///
+/// Adds a !type annotation for every trait in the trait hierarchy, and a !vcall_visibility
+/// annotation based on the most visible trait in the hierarchy.
 pub(crate) fn apply_vcall_visibility_metadata<'ll, 'tcx>(
     cx: &CodegenCx<'ll, 'tcx>,
-    ty: Ty<'tcx>,
-    trait_ref: Option<ExistentialTraitRef<'tcx>>,
+    trait_hierarchy: &[ExistentialTraitRef<'tcx>],
     vtable: &'ll Value,
 ) {
     // FIXME(flip1995): The virtual function elimination optimization only works with full LTO in
@@ -1745,51 +1791,30 @@ pub(crate) fn apply_vcall_visibility_metadata<'ll, 'tcx>(
         return;
     }
 
-    enum VCallVisibility {
-        Public = 0,
-        LinkageUnit = 1,
-        TranslationUnit = 2,
-    }
-
-    let Some(trait_ref) = trait_ref else { return };
+    let Some(vcall_visibility) = trait_hierarchy
+        .iter()
+        .map(|trait_ref| VCallVisibility::for_def_id(cx, trait_ref.def_id))
+        .min()
+    else {
+        return;
+    };
 
-    // Unwrap potential addrspacecast
+    // Unwrap a potential addrspacecast.
     let vtable = find_vtable_behind_cast(vtable);
-    let trait_ref_self = trait_ref.with_self_ty(cx.tcx, ty);
-    let trait_def_id = trait_ref_self.def_id;
-    let trait_vis = cx.tcx.visibility(trait_def_id);
-
-    let cgus = cx.sess().codegen_units().as_usize();
-    let single_cgu = cgus == 1;
-
-    let lto = cx.sess().lto();
-
-    // Since LLVM requires full LTO for the virtual function elimination optimization to apply,
-    // only the `Lto::Fat` cases are relevant currently.
-    let vcall_visibility = match (lto, trait_vis, single_cgu) {
-        // If there is not LTO and the visibility in public, we have to assume that the vtable can
-        // be seen from anywhere. With multiple CGUs, the vtable is quasi-public.
-        (Lto::No | Lto::ThinLocal, Visibility::Public, _)
-        | (Lto::No, Visibility::Restricted(_), false) => VCallVisibility::Public,
-        // With LTO and a quasi-public visibility, the usages of the functions of the vtable are
-        // all known by the `LinkageUnit`.
-        // FIXME: LLVM only supports this optimization for `Lto::Fat` currently. Once it also
-        // supports `Lto::Thin` the `VCallVisibility` may have to be adjusted for those.
-        (Lto::Fat | Lto::Thin, Visibility::Public, _)
-        | (Lto::ThinLocal | Lto::Thin | Lto::Fat, Visibility::Restricted(_), false) => {
-            VCallVisibility::LinkageUnit
-        }
-        // If there is only one CGU, private vtables can only be seen by that CGU/translation unit
-        // and therefore we know of all usages of functions in the vtable.
-        (_, Visibility::Restricted(_), true) => VCallVisibility::TranslationUnit,
-    };
 
-    let trait_ref_typeid = typeid_for_trait_ref(cx.tcx, trait_ref);
-    let typeid = cx.create_metadata(trait_ref_typeid.as_bytes());
+    // The start of a Rust vtable is at offset 0 of a `&dyn T`'s address.
+    let vtable_offset = 0;
 
-    let type_ = [llvm::LLVMValueAsMetadata(cx.const_usize(0)), typeid];
-    cx.global_add_metadata_node(vtable, llvm::MD_type, &type_);
+    // Add `!type !{i64 0, !"typeid"}` to the vtable for each trait in the hierarchy.
+    for trait_ref in trait_hierarchy {
+        let trait_ref_typeid = typeid_for_trait_ref(cx.tcx, *trait_ref);
+        let typeid = cx.create_metadata(trait_ref_typeid.as_bytes());
+
+        let type_ = [llvm::LLVMValueAsMetadata(cx.const_u64(vtable_offset)), typeid];
+        cx.global_add_metadata_node(vtable, llvm::MD_type, &type_);
+    }
 
+    // Add `!vcall_visibility !{i64 vis}` to the vtable.
     let vcall_visibility = [llvm::LLVMValueAsMetadata(cx.const_u64(vcall_visibility as u64))];
     cx.global_set_metadata_node(vtable, llvm::MD_vcall_visibility, &vcall_visibility);
 }

compiler/rustc_codegen_ssa/src/meth.rs

@@ -3,6 +3,7 @@ use rustc_session::config::Lto;
 use rustc_span::bug;
 use rustc_symbol_mangling::typeid_for_trait_ref;
 use rustc_target::callconv::FnAbi;
+use rustc_trait_selection::traits::vtable;
 use tracing::{debug, instrument};
 
 use crate::traits::*;
@@ -143,22 +144,32 @@ fn dyn_trait_in_self<'tcx>(
 #[instrument(level = "debug", skip(cx))]
 pub(crate) fn get_vtable<'tcx, Cx: CodegenMethods<'tcx>>(
     cx: &Cx,
-    ty: Ty<'tcx>,
+    self_ty: Ty<'tcx>,
     trait_ref: Option<ty::ExistentialTraitRef<'tcx>>,
 ) -> Cx::Value {
     let tcx = cx.tcx();
 
     // Check the cache.
-    if let Some(&val) = cx.vtables().borrow().get(&(ty, trait_ref)) {
+    if let Some(&val) = cx.vtables().borrow().get(&(self_ty, trait_ref)) {
         return val;
     }
 
-    let vtable_alloc_id = tcx.vtable_allocation((ty, trait_ref));
+    let vtable_alloc_id = tcx.vtable_allocation((self_ty, trait_ref));
     let vtable_allocation = tcx.global_alloc(vtable_alloc_id).unwrap_memory();
     let vtable = cx.static_addr_of(vtable_allocation, Some("vtable"));
 
-    cx.apply_vcall_visibility_metadata(ty, trait_ref, vtable);
-    cx.create_vtable_debuginfo(ty, trait_ref, vtable);
-    cx.vtables().borrow_mut().insert((ty, trait_ref), vtable);
+    if let Some(trait_ref) = trait_ref
+        && cx.sess().opts.unstable_opts.virtual_function_elimination
+        && cx.sess().lto() == Lto::Fat
+    {
+        let key = trait_ref.with_self_ty(tcx, self_ty);
+        let trait_hierarchy = vtable::vtable_prefix_traits(tcx, key);
+
+        cx.apply_vcall_visibility_metadata(&trait_hierarchy, vtable);
+    }
+
+    cx.create_vtable_debuginfo(self_ty, trait_ref, vtable);
+    cx.vtables().borrow_mut().insert((self_ty, trait_ref), vtable);
+
     vtable
 }

compiler/rustc_codegen_ssa/src/traits/misc.rs

@@ -13,8 +13,7 @@ pub trait MiscCodegenMethods<'tcx>: BackendTypes {
     ) -> &RefCell<FxHashMap<(Ty<'tcx>, Option<ty::ExistentialTraitRef<'tcx>>), Self::Value>>;
     fn apply_vcall_visibility_metadata(
         &self,
-        _ty: Ty<'tcx>,
-        _poly_trait_ref: Option<ty::ExistentialTraitRef<'tcx>>,
+        _trait_hierarchy: &[ty::ExistentialTraitRef<'tcx>],
         _vtable: Self::Value,
     ) {
     }

compiler/rustc_trait_selection/src/traits/vtable.rs

@@ -16,7 +16,9 @@ use crate::traits::is_vtable_safe_method;
 
 #[derive(Clone, Debug)]
 pub enum VtblSegment<'tcx> {
+    /// The segment containing Drop, Size and Align.
     MetadataDSA,
+    /// A segment with the methods of a particular trait.
     TraitOwnEntries { trait_ref: ty::TraitRef<'tcx>, emit_vptr: bool },
 }
 
@@ -453,6 +455,38 @@ pub(crate) fn supertrait_vtable_slot<'tcx>(
     })
 }
 
+/// Returns the traits that are prefixes of this vtable, such that upcasting to these traits
+/// reuses the vtable pointer unchanged.
+///
+/// For a `trait Foo: A + B`, the methods for A are stored as a prefix to the vtable for Foo.
+/// If B has no methods, it is also part of this prefix (although it takes up no space).
+///
+/// This function returns the full set of these traits (A might have its own supertrait, etc.).
+/// The input trait is the last element of the output vector.
+pub fn vtable_prefix_traits<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    key: ty::TraitRef<'tcx>,
+) -> Vec<ty::ExistentialTraitRef<'tcx>> {
+    debug_assert!(!key.has_non_region_infer() && !key.has_non_region_param());
+
+    let mut entries = vec![];
+
+    let vtable_segment_callback = |segment| -> ControlFlow<()> {
+        if let VtblSegment::TraitOwnEntries { trait_ref, emit_vptr: false } = segment {
+            entries.push(ty::ExistentialTraitRef::erase_self_ty(tcx, trait_ref));
+        }
+
+        ControlFlow::Continue(())
+    };
+
+    let _ = prepare_vtable_segments(tcx, key, vtable_segment_callback);
+
+    // The input trait is the last element in the output vector.
+    debug_assert_eq!(entries.last(), Some(&ty::ExistentialTraitRef::erase_self_ty(tcx, key)));
+
+    entries
+}
+
 pub(super) fn provide(providers: &mut Providers) {
     *providers = Providers {
         own_existential_vtable_entries,

tests/codegen-llvm/vfe-annotate-supertrait.rs

@@ -0,0 +1,95 @@
+//@ compile-flags: -Zvirtual-function-elimination -Clto=fat -Copt-level=3 -Zverify-llvm-ir
+//@ no-prefer-dynamic
+//@ min-llvm-version: 23
+#![crate_type = "staticlib"]
+
+// Check that vtables include a !type annotation for supertraits whose vtable are a prefix.
+//
+// What is still missing is that for any supertrait whose vtable is not a prefix, its vtable does
+// not get any !type annotation. That means we provide incomplete information to VFE and it may
+// incorrectly replace a virtual call with a direct call to the only callee VFE knows about.
+
+use std::hint::black_box;
+
+// FIXME: C_VTABLE should have a !type annotation for C, but it currently has none.
+// CHECK: @vtable.{{[0-9]+}} = {{.*}}3Foo3foo{{.*}} !type ![[META_SIZED:[0-9]+]], !type ![[A:[0-9]+]], !type ![[B:[0-9]+]], !type ![[FOO:[0-9]+]], !guid
+// CHECK: @[[C_VTABLE:alloc_[0-9a-f]+]] = {{.*}}1C1c }>, align {{[0-9]+}}, !guid
+// CHECK: @vtable.{{[0-9]+}} = {{.*}}1C1c, ptr @[[C_VTABLE]], ptr {{.*}}3Bar3bar{{.*}} !type ![[META_SIZED]], !type ![[A]], !type ![[BAR:[0-9]+]], !guid
+
+// CHECK-DAG: ![[META_SIZED]] = !{i64 0, !"{{.*}}9MetaSized"}
+// CHECK-DAG: ![[A]] = !{i64 0, !"{{.*}}1A"}
+// CHECK-DAG: ![[B]] = !{i64 0, !"{{.*}}1B"}
+// CHECK-DAG: ![[FOO]] = !{i64 0, !"{{.*}}3Foo"}
+// CHECK-DAG: ![[BAR]] = !{i64 0, !"{{.*}}3Bar"}
+
+pub trait A {
+    fn a(&self) -> i32;
+}
+
+pub trait B {}
+
+pub trait C {
+    fn c(&self) -> i32;
+}
+
+/// A trait with two supertraits:
+///
+/// - A is a prefix of Foo, because it is the first supertrait in the list
+/// - B is a prefix of Foo, because it has no methods
+pub trait Foo: A + B {
+    fn foo(&self) -> i32;
+}
+
+/// A trait with two supertraits:
+///
+/// - A is a prefix of Bar, because it is the first supertrait in the list
+/// - C is a NOT a prefix of Bar, because it has methods
+pub trait Bar: A + C {
+    fn bar(&self) -> i32;
+}
+
+// The remainder is uninteresting, and only here so that the optimizer can't optimize out the
+// interesting bits.
+
+pub struct S;
+
+impl A for S {
+    fn a(&self) -> i32 {
+        black_box(1)
+    }
+}
+
+impl B for S {}
+
+impl C for S {
+    fn c(&self) -> i32 {
+        black_box(2)
+    }
+}
+
+impl Foo for S {
+    fn foo(&self) -> i32 {
+        black_box(3)
+    }
+}
+
+impl Bar for S {
+    fn bar(&self) -> i32 {
+        black_box(4)
+    }
+}
+
+#[unsafe(no_mangle)]
+pub fn use_foo(x: &dyn Foo) -> i32 {
+    x.foo() + x.a()
+}
+
+#[unsafe(no_mangle)]
+pub fn use_bar(x: &dyn Bar) -> i32 {
+    x.bar() + x.a() + x.c()
+}
+
+#[unsafe(no_mangle)]
+pub fn use_s() -> i32 {
+    use_foo(black_box(&S)) + use_bar(black_box(&S))
+}