CFI: Fix many vtable-related problems

compiler/rustc_codegen_llvm/src/callee.rs

@@ -67,14 +67,13 @@ pub fn get_fn<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>, instance: Instance<'tcx>) ->
                     true,
                 ),
                 fn_abi,
-                Some(instance),
             );
             unsafe {
                 llvm::LLVMSetDLLStorageClass(llfn, llvm::DLLStorageClass::DllImport);
             }
             llfn
         } else {
-            cx.declare_fn(sym, fn_abi, Some(instance))
+            cx.declare_fn(sym, fn_abi)
         };
         debug!("get_fn: not casting pointer!");
 

compiler/rustc_codegen_llvm/src/declare.rs

@@ -19,11 +19,8 @@ use crate::llvm::AttributePlace::Function;
 use crate::type_::Type;
 use crate::value::Value;
 use rustc_codegen_ssa::traits::TypeMembershipMethods;
-use rustc_middle::ty::{Instance, Ty};
-use rustc_symbol_mangling::typeid::{
-    kcfi_typeid_for_fnabi, kcfi_typeid_for_instance, typeid_for_fnabi, typeid_for_instance,
-    TypeIdOptions,
-};
+use rustc_middle::ty::Ty;
+use rustc_symbol_mangling::typeid::{kcfi_typeid_for_fnabi, typeid_for_fnabi, TypeIdOptions};
 use smallvec::SmallVec;
 
 /// Declare a function.
@@ -119,12 +116,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
     ///
     /// If there’s a value with the same name already declared, the function will
     /// update the declaration and return existing Value instead.
-    pub fn declare_fn(
-        &self,
-        name: &str,
-        fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
-        instance: Option<Instance<'tcx>>,
-    ) -> &'ll Value {
+    pub fn declare_fn(&self, name: &str, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> &'ll Value {
         debug!("declare_rust_fn(name={:?}, fn_abi={:?})", name, fn_abi);
 
         // Function addresses in Rust are never significant, allowing functions to
@@ -140,35 +132,18 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
         fn_abi.apply_attrs_llfn(self, llfn);
 
         if self.tcx.sess.is_sanitizer_cfi_enabled() {
-            if let Some(instance) = instance {
-                let typeid = typeid_for_instance(self.tcx, &instance, TypeIdOptions::empty());
-                self.set_type_metadata(llfn, typeid);
-                let typeid =
-                    typeid_for_instance(self.tcx, &instance, TypeIdOptions::GENERALIZE_POINTERS);
-                self.add_type_metadata(llfn, typeid);
-                let typeid =
-                    typeid_for_instance(self.tcx, &instance, TypeIdOptions::NORMALIZE_INTEGERS);
-                self.add_type_metadata(llfn, typeid);
-                let typeid = typeid_for_instance(
-                    self.tcx,
-                    &instance,
-                    TypeIdOptions::GENERALIZE_POINTERS | TypeIdOptions::NORMALIZE_INTEGERS,
-                );
-                self.add_type_metadata(llfn, typeid);
-            } else {
-                let typeid = typeid_for_fnabi(self.tcx, fn_abi, TypeIdOptions::empty());
-                self.set_type_metadata(llfn, typeid);
-                let typeid = typeid_for_fnabi(self.tcx, fn_abi, TypeIdOptions::GENERALIZE_POINTERS);
-                self.add_type_metadata(llfn, typeid);
-                let typeid = typeid_for_fnabi(self.tcx, fn_abi, TypeIdOptions::NORMALIZE_INTEGERS);
-                self.add_type_metadata(llfn, typeid);
-                let typeid = typeid_for_fnabi(
-                    self.tcx,
-                    fn_abi,
-                    TypeIdOptions::GENERALIZE_POINTERS | TypeIdOptions::NORMALIZE_INTEGERS,
-                );
-                self.add_type_metadata(llfn, typeid);
-            }
+            let typeid = typeid_for_fnabi(self.tcx, fn_abi, TypeIdOptions::empty());
+            self.set_type_metadata(llfn, typeid);
+            let typeid = typeid_for_fnabi(self.tcx, fn_abi, TypeIdOptions::GENERALIZE_POINTERS);
+            self.add_type_metadata(llfn, typeid);
+            let typeid = typeid_for_fnabi(self.tcx, fn_abi, TypeIdOptions::NORMALIZE_INTEGERS);
+            self.add_type_metadata(llfn, typeid);
+            let typeid = typeid_for_fnabi(
+                self.tcx,
+                fn_abi,
+                TypeIdOptions::GENERALIZE_POINTERS | TypeIdOptions::NORMALIZE_INTEGERS,
+            );
+            self.add_type_metadata(llfn, typeid);
         }
 
         if self.tcx.sess.is_sanitizer_kcfi_enabled() {
@@ -181,13 +156,8 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
                 options.insert(TypeIdOptions::NORMALIZE_INTEGERS);
             }
 
-            if let Some(instance) = instance {
-                let kcfi_typeid = kcfi_typeid_for_instance(self.tcx, &instance, options);
-                self.set_kcfi_type_metadata(llfn, kcfi_typeid);
-            } else {
-                let kcfi_typeid = kcfi_typeid_for_fnabi(self.tcx, fn_abi, options);
-                self.set_kcfi_type_metadata(llfn, kcfi_typeid);
-            }
+            let kcfi_typeid = kcfi_typeid_for_fnabi(self.tcx, fn_abi, options);
+            self.set_kcfi_type_metadata(llfn, kcfi_typeid);
         }
 
         llfn

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -939,7 +939,7 @@ fn gen_fn<'ll, 'tcx>(
 ) -> (&'ll Type, &'ll Value) {
     let fn_abi = cx.fn_abi_of_fn_ptr(rust_fn_sig, ty::List::empty());
     let llty = fn_abi.llvm_type(cx);
-    let llfn = cx.declare_fn(name, fn_abi, None);
+    let llfn = cx.declare_fn(name, fn_abi);
     cx.set_frame_pointer_type(llfn);
     cx.apply_target_cpu_attr(llfn);
     // FIXME(eddyb) find a nicer way to do this.

compiler/rustc_codegen_llvm/src/mono_item.rs

@@ -51,7 +51,7 @@ impl<'tcx> PreDefineMethods<'tcx> for CodegenCx<'_, 'tcx> {
         assert!(!instance.args.has_infer());
 
         let fn_abi = self.fn_abi_of_instance(instance, ty::List::empty());
-        let lldecl = self.declare_fn(symbol_name, fn_abi, Some(instance));
+        let lldecl = self.declare_fn(symbol_name, fn_abi);
         unsafe { llvm::LLVMRustSetLinkage(lldecl, base::linkage_to_llvm(linkage)) };
         let attrs = self.tcx.codegen_fn_attrs(instance.def_id());
         base::set_link_section(lldecl, attrs);

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -495,7 +495,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     //
                     let virtual_drop = Instance {
                         def: ty::InstanceDef::Virtual(drop_fn.def_id(), 0),
-                        args: drop_fn.args,
+                        args: bx.tcx().strip_receiver_auto(drop_fn.args),
                     };
                     debug!("ty = {:?}", ty);
                     debug!("drop_fn = {:?}", drop_fn);
@@ -535,7 +535,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     // SO THEN WE CAN USE THE ABOVE CODE.
                     let virtual_drop = Instance {
                         def: ty::InstanceDef::Virtual(drop_fn.def_id(), 0),
-                        args: drop_fn.args,
+                        args: bx.tcx().strip_receiver_auto(drop_fn.args),
                     };
                     debug!("ty = {:?}", ty);
                     debug!("drop_fn = {:?}", drop_fn);

compiler/rustc_const_eval/src/interpret/terminator.rs

@@ -547,6 +547,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             | ty::InstanceDef::CloneShim(..)
             | ty::InstanceDef::FnPtrAddrShim(..)
             | ty::InstanceDef::ThreadLocalShim(..)
+            | ty::InstanceDef::CfiShim { .. }
             | ty::InstanceDef::Item(_) => {
                 // We need MIR for this fn
                 let Some((body, instance)) = M::find_mir_or_eval_fn(

compiler/rustc_middle/src/mir/mono.rs

@@ -406,7 +406,8 @@ impl<'tcx> CodegenUnit<'tcx> {
                             | InstanceDef::DropGlue(..)
                             | InstanceDef::CloneShim(..)
                             | InstanceDef::ThreadLocalShim(..)
-                            | InstanceDef::FnPtrAddrShim(..) => None,
+                            | InstanceDef::FnPtrAddrShim(..)
+                            | InstanceDef::CfiShim { .. } => None,
                         }
                     }
                     MonoItem::Static(def_id) => def_id.as_local().map(Idx::index),

compiler/rustc_middle/src/mir/visit.rs

@@ -331,33 +331,10 @@ macro_rules! make_mir_visitor {
                     self.visit_source_scope($(& $mutability)? *parent_scope);
                 }
                 if let Some((callee, callsite_span)) = inlined {
-                    let location = Location::START;
-
                     self.visit_span($(& $mutability)? *callsite_span);
-
-                    let ty::Instance { def: callee_def, args: callee_args } = callee;
-                    match callee_def {
-                        ty::InstanceDef::Item(_def_id) => {}
-
-                        ty::InstanceDef::Intrinsic(_def_id) |
-                        ty::InstanceDef::VTableShim(_def_id) |
-                        ty::InstanceDef::ReifyShim(_def_id) |
-                        ty::InstanceDef::Virtual(_def_id, _) |
-                        ty::InstanceDef::ThreadLocalShim(_def_id) |
-                        ty::InstanceDef::ClosureOnceShim { call_once: _def_id, track_caller: _ } |
-                        ty::InstanceDef::ConstructCoroutineInClosureShim { coroutine_closure_def_id: _def_id, target_kind: _ } |
-                        ty::InstanceDef::CoroutineKindShim { coroutine_def_id: _def_id, target_kind: _ } |
-                        ty::InstanceDef::DropGlue(_def_id, None) => {}
-
-                        ty::InstanceDef::FnPtrShim(_def_id, ty) |
-                        ty::InstanceDef::DropGlue(_def_id, Some(ty)) |
-                        ty::InstanceDef::CloneShim(_def_id, ty) |
-                        ty::InstanceDef::FnPtrAddrShim(_def_id, ty) => {
-                            // FIXME(eddyb) use a better `TyContext` here.
-                            self.visit_ty($(& $mutability)? *ty, TyContext::Location(location));
-                        }
-                    }
-                    self.visit_args(callee_args, location);
+                    let ty::Instance { def, args } = callee;
+                    self.visit_instance_def($(& $mutability)? *def);
+                    self.visit_args(& $($mutability)? *args, Location::START);
                 }
                 if let Some(inlined_parent_scope) = inlined_parent_scope {
                     self.visit_source_scope($(& $mutability)? *inlined_parent_scope);
@@ -940,6 +917,46 @@ macro_rules! make_mir_visitor {
 
             // Convenience methods
 
+            fn visit_instance_def(&mut self, def: $(& $mutability)? ty::InstanceDef<'tcx>) {
+                let location = Location::START;
+                match def {
+                    ty::InstanceDef::Item(_def_id) => {}
+
+                    ty::InstanceDef::Intrinsic(_def_id) |
+                    ty::InstanceDef::VTableShim(_def_id) |
+                    ty::InstanceDef::ReifyShim(_def_id) |
+                    ty::InstanceDef::Virtual(_def_id, _) |
+                    ty::InstanceDef::ThreadLocalShim(_def_id) |
+                    ty::InstanceDef::ClosureOnceShim { call_once: _def_id, track_caller: _ } |
+                    ty::InstanceDef::ConstructCoroutineInClosureShim { coroutine_closure_def_id: _def_id, target_kind: _ } |
+                    ty::InstanceDef::CoroutineKindShim { coroutine_def_id: _def_id, target_kind: _ } |
+                    ty::InstanceDef::DropGlue(_def_id, None) => {}
+
+                    ty::InstanceDef::FnPtrShim(_def_id, ty) |
+                    ty::InstanceDef::DropGlue(_def_id, Some(ty)) |
+                    ty::InstanceDef::CloneShim(_def_id, ty) |
+                    ty::InstanceDef::FnPtrAddrShim(_def_id, ty) => {
+                        // FIXME(eddyb) use a better `TyContext` here.
+                        self.visit_ty($(& $mutability *)? ty, TyContext::Location(location));
+                    }
+                    ty::InstanceDef::CfiShim { target_instance, invoke_ty } => {
+                        self.visit_ty($(& $mutability *)? invoke_ty, TyContext::Location(location));
+                        let $($mutability)? local_target_instance = {
+                            $(
+                                let $mutability _unused = ();
+                                *
+                            )?
+                            *target_instance
+                        };
+                        self.visit_instance_def($(& $mutability)? local_target_instance);
+                        $(
+                            *target_instance = self.tcx().arena.alloc(local_target_instance);
+                            let $mutability _unused = ();
+                        )?
+                    }
+                }
+            }
+
             fn visit_location(
                 &mut self,
                 body: &$($mutability)? Body<'tcx>,

compiler/rustc_middle/src/query/mod.rs

@@ -2247,6 +2247,22 @@ rustc_queries! {
     query find_field((def_id, ident): (DefId, rustc_span::symbol::Ident)) -> Option<rustc_target::abi::FieldIdx> {
         desc { |tcx| "find the index of maybe nested field `{ident}` in `{}`", tcx.def_path_str(def_id) }
     }
+
+    /// Construct a type for a trait object corresponding to `trait_ref`. This type will have all
+    /// associated types for it and its supertraits expanded and resolved as additional predicates.
+    ///
+    /// The provided `trait_ref` must be sufficiently instantiated that all associated types can be
+    /// successfully resolved.
+    query trait_object_ty(trait_ref: ty::PolyTraitRef<'tcx>) -> Ty<'tcx> {
+        desc { "Compute the trait object type for calling a method on a trait" }
+    }
+
+    /// Strip auto traits off the first parameter in the parametr list. Intended for use when
+    /// constructing `InstanceDef::Virtual`, as auto traits won't be part of the vtable's `Self`
+    /// types.
+    query strip_receiver_auto(args: ty::GenericArgsRef<'tcx>) -> ty::GenericArgsRef<'tcx> {
+        desc { "Strip auto traits off the first type arg" }
+    }
 }
 
 rustc_query_append! { define_callbacks! }

compiler/rustc_middle/src/ty/codec.rs

@@ -526,6 +526,7 @@ impl_arena_copy_decoder! {<'tcx>
     rustc_span::def_id::LocalDefId,
     (rustc_middle::middle::exported_symbols::ExportedSymbol<'tcx>, rustc_middle::middle::exported_symbols::SymbolExportInfo),
     ty::DeducedParamAttrs,
+    ty::InstanceDef<'tcx>,
 }
 
 #[macro_export]