Add llvm.type.checked.load intrinsic

Add the intrinsic

declare {i8*, i1} @llvm.type.checked.load(i8* %ptr, i32 %offset, metadata %type)

This is used in the VFE optimization when lowering loading functions
from vtables to LLVM IR. The `metadata` is used to map the function to
all vtables this function could belong to. This ensures that functions
from vtables that might be used somewhere won't get removed.

6 files changed, 88 insertions, 18 deletions

diff --git a/compiler/rustc_codegen_gcc/src/intrinsic/mod.rs b/compiler/rustc_codegen_gcc/src/intrinsic/mod.rs
index c6681de68e2..5fbdedac0c4 100644
--- a/compiler/rustc_codegen_gcc/src/intrinsic/mod.rs
+++ b/compiler/rustc_codegen_gcc/src/intrinsic/mod.rs
@@ -356,6 +356,16 @@ impl<'a, 'gcc, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'a, 'gcc, 'tcx> {
         self.context.new_rvalue_from_int(self.int_type, 0)
     }
 
+    fn type_checked_load(
+        &mut self,
+        _llvtable: Self::Value,
+        _vtable_byte_offset: u64,
+        _typeid: Self::Value,
+    ) -> Self::Value {
+        // Unsupported.
+        self.context.new_rvalue_from_int(self.int_type, 0)
+    }
+
     fn va_start(&mut self, _va_list: RValue<'gcc>) -> RValue<'gcc> {
         unimplemented!();
     }
diff --git a/compiler/rustc_codegen_llvm/src/context.rs b/compiler/rustc_codegen_llvm/src/context.rs
index 27fe60161d5..b5c31fcebe0 100644
--- a/compiler/rustc_codegen_llvm/src/context.rs
+++ b/compiler/rustc_codegen_llvm/src/context.rs
@@ -665,6 +665,7 @@ impl<'ll> CodegenCx<'ll, '_> {
         let t_isize = self.type_isize();
         let t_f32 = self.type_f32();
         let t_f64 = self.type_f64();
+        let t_metadata = self.type_metadata();
 
         ifn!("llvm.wasm.trunc.unsigned.i32.f32", fn(t_f32) -> t_i32);
         ifn!("llvm.wasm.trunc.unsigned.i32.f64", fn(t_f64) -> t_i32);
@@ -890,11 +891,12 @@ impl<'ll> CodegenCx<'ll, '_> {
             ifn!("llvm.instrprof.increment", fn(i8p, t_i64, t_i32, t_i32) -> void);
         }
 
-        ifn!("llvm.type.test", fn(i8p, self.type_metadata()) -> i1);
+        ifn!("llvm.type.test", fn(i8p, t_metadata) -> i1);
+        ifn!("llvm.type.checked.load", fn(i8p, t_i32, t_metadata) -> mk_struct! {i8p, i1});
 
         if self.sess().opts.debuginfo != DebugInfo::None {
-            ifn!("llvm.dbg.declare", fn(self.type_metadata(), self.type_metadata()) -> void);
-            ifn!("llvm.dbg.value", fn(self.type_metadata(), t_i64, self.type_metadata()) -> void);
+            ifn!("llvm.dbg.declare", fn(t_metadata, t_metadata) -> void);
+            ifn!("llvm.dbg.value", fn(t_metadata, t_i64, t_metadata) -> void);
         }
         None
     }
diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs
index 9927f5f399b..a18f5b9dd7f 100644
--- a/compiler/rustc_codegen_llvm/src/intrinsic.rs
+++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs
@@ -406,6 +406,16 @@ impl<'ll, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'_, 'll, 'tcx> {
         self.call_intrinsic("llvm.type.test", &[bitcast, typeid])
     }
 
+    fn type_checked_load(
+        &mut self,
+        llvtable: &'ll Value,
+        vtable_byte_offset: u64,
+        typeid: &'ll Value,
+    ) -> Self::Value {
+        let vtable_byte_offset = self.const_i32(vtable_byte_offset as i32);
+        self.call_intrinsic("llvm.type.checked.load", &[llvtable, vtable_byte_offset, typeid])
+    }
+
     fn va_start(&mut self, va_list: &'ll Value) -> &'ll Value {
         self.call_intrinsic("llvm.va_start", &[va_list])
     }
diff --git a/compiler/rustc_codegen_ssa/src/meth.rs b/compiler/rustc_codegen_ssa/src/meth.rs
index 00f101595f2..5203ebfad75 100644
--- a/compiler/rustc_codegen_ssa/src/meth.rs
+++ b/compiler/rustc_codegen_ssa/src/meth.rs
@@ -1,6 +1,8 @@
 use crate::traits::*;
 
-use rustc_middle::ty::{self, Ty};
+use rustc_middle::ty::{self, subst::GenericArgKind, ExistentialPredicate, Ty, TyCtxt};
+use rustc_session::config::Lto;
+use rustc_symbol_mangling::typeid_for_trait_ref;
 use rustc_target::abi::call::FnAbi;
 
 #[derive(Copy, Clone, Debug)]
@@ -15,20 +17,32 @@ impl<'a, 'tcx> VirtualIndex {
         self,
         bx: &mut Bx,
         llvtable: Bx::Value,
+        ty: Ty<'tcx>,
         fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
     ) -> Bx::Value {
         // Load the data pointer from the object.
-        debug!("get_fn({:?}, {:?})", llvtable, self);
-
+        debug!("get_fn({llvtable:?}, {ty:?}, {self:?})");
         let llty = bx.fn_ptr_backend_type(fn_abi);
         let llvtable = bx.pointercast(llvtable, bx.type_ptr_to(llty));
-        let ptr_align = bx.tcx().data_layout.pointer_align.abi;
-        let gep = bx.inbounds_gep(llty, llvtable, &[bx.const_usize(self.0)]);
-        let ptr = bx.load(llty, gep, ptr_align);
-        bx.nonnull_metadata(ptr);
-        // Vtable loads are invariant.
-        bx.set_invariant_load(ptr);
-        ptr
+
+        if bx.cx().sess().opts.debugging_opts.virtual_function_elimination
+            && bx.cx().sess().lto() == Lto::Fat
+        {
+            let typeid =
+                bx.typeid_metadata(typeid_for_trait_ref(bx.tcx(), get_trait_ref(bx.tcx(), ty)));
+            let vtable_byte_offset = self.0 * bx.data_layout().pointer_size.bytes();
+            let type_checked_load = bx.type_checked_load(llvtable, vtable_byte_offset, typeid);
+            let func = bx.extract_value(type_checked_load, 0);
+            bx.pointercast(func, llty)
+        } else {
+            let ptr_align = bx.tcx().data_layout.pointer_align.abi;
+            let gep = bx.inbounds_gep(llty, llvtable, &[bx.const_usize(self.0)]);
+            let ptr = bx.load(llty, gep, ptr_align);
+            bx.nonnull_metadata(ptr);
+            // Vtable loads are invariant.
+            bx.set_invariant_load(ptr);
+            ptr
+        }
     }
 
     pub fn get_usize<Bx: BuilderMethods<'a, 'tcx>>(
@@ -50,6 +64,24 @@ impl<'a, 'tcx> VirtualIndex {
     }
 }
 
+fn get_trait_ref<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> ty::PolyExistentialTraitRef<'tcx> {
+    for arg in ty.peel_refs().walk() {
+        if let GenericArgKind::Type(ty) = arg.unpack() {
+            if let ty::Dynamic(trait_refs, _) = ty.kind() {
+                return trait_refs[0].map_bound(|trait_ref| match trait_ref {
+                    ExistentialPredicate::Trait(tr) => tr,
+                    ExistentialPredicate::Projection(proj) => proj.trait_ref(tcx),
+                    ExistentialPredicate::AutoTrait(_) => {
+                        bug!("auto traits don't have functions")
+                    }
+                });
+            }
+        }
+    }
+
+    bug!("expected a `dyn Trait` ty, found {ty:?}")
+}
+
 /// Creates a dynamic vtable for the given type and vtable origin.
 /// This is used only for objects.
 ///
diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
index 03ef6d50d44..db348f2bdd5 100644
--- a/compiler/rustc_codegen_ssa/src/mir/block.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/block.rs
@@ -401,7 +401,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 args = &args[..1];
                 (
                     meth::VirtualIndex::from_index(ty::COMMON_VTABLE_ENTRIES_DROPINPLACE)
-                        .get_fn(&mut bx, vtable, &fn_abi),
+                        .get_fn(&mut bx, vtable, ty, &fn_abi),
                     fn_abi,
                 )
             }
@@ -819,9 +819,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     // the data pointer as the first argument
                     match op.val {
                         Pair(data_ptr, meta) => {
-                            llfn = Some(
-                                meth::VirtualIndex::from_index(idx).get_fn(&mut bx, meta, &fn_abi),
-                            );
+                            llfn = Some(meth::VirtualIndex::from_index(idx).get_fn(
+                                &mut bx,
+                                meta,
+                                op.layout.ty,
+                                &fn_abi,
+                            ));
                             llargs.push(data_ptr);
                             continue 'make_args;
                         }
@@ -829,7 +832,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     }
                 } else if let Ref(data_ptr, Some(meta), _) = op.val {
                     // by-value dynamic dispatch
-                    llfn = Some(meth::VirtualIndex::from_index(idx).get_fn(&mut bx, meta, &fn_abi));
+                    llfn = Some(meth::VirtualIndex::from_index(idx).get_fn(
+                        &mut bx,
+                        meta,
+                        op.layout.ty,
+                        &fn_abi,
+                    ));
                     llargs.push(data_ptr);
                     continue;
                 } else {
diff --git a/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs b/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs
index 02be6cd360c..7755e67938c 100644
--- a/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs
@@ -22,6 +22,14 @@ pub trait IntrinsicCallMethods<'tcx>: BackendTypes {
     fn expect(&mut self, cond: Self::Value, expected: bool) -> Self::Value;
     /// Trait method used to test whether a given pointer is associated with a type identifier.
     fn type_test(&mut self, pointer: Self::Value, typeid: Self::Value) -> Self::Value;
+    /// Trait method used to load a function while testing if it is associated with a type
+    /// identifier.
+    fn type_checked_load(
+        &mut self,
+        llvtable: Self::Value,
+        vtable_byte_offset: u64,
+        typeid: Self::Value,
+    ) -> Self::Value;
     /// Trait method used to inject `va_start` on the "spoofed" `VaListImpl` in
     /// Rust defined C-variadic functions.
     fn va_start(&mut self, val: Self::Value) -> Self::Value;