CFI: Support arbitrary receivers

Previously, we only rewrote `&self` and `&mut self` receivers. By
instantiating the method from the trait definition, we can make this
work work with arbitrary legal receivers instead.

compiler/rustc_symbol_mangling/src/lib.rs

@@ -90,6 +90,7 @@
 #![doc(html_root_url = "https://doc.rust-lang.org/nightly/nightly-rustc/")]
 #![doc(rust_logo)]
 #![feature(rustdoc_internals)]
+#![feature(let_chains)]
 #![allow(internal_features)]
 
 #[macro_use]

compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs

@@ -1115,51 +1115,31 @@ pub fn typeid_for_instance<'tcx>(
         instance.args = strip_receiver_auto(tcx, instance.args)
     }
 
+    if let Some(impl_id) = tcx.impl_of_method(instance.def_id())
+        && let Some(trait_ref) = tcx.impl_trait_ref(impl_id)
+    {
+        // Trait methods will have a Self polymorphic parameter, where the concreteized
+        // implementatation will not. We need to walk back to the more general trait method
+        let trait_ref = trait_ref.instantiate(tcx, instance.args);
+        let invoke_ty = tcx.trait_object_ty(ty::Binder::dummy(trait_ref));
+        let method_id = tcx
+            .impl_item_implementor_ids(impl_id)
+            .items()
+            .filter_map(|(trait_method, impl_method)| {
+                (*impl_method == instance.def_id()).then_some(*trait_method)
+            })
+            .min()
+            .unwrap();
+        instance.def = ty::InstanceDef::Virtual(method_id, 0);
+        instance.args = tcx.mk_args_trait(invoke_ty, trait_ref.args.into_iter().skip(1));
+    }
+
     let fn_abi = tcx
         .fn_abi_of_instance(tcx.param_env(instance.def_id()).and((instance, ty::List::empty())))
         .unwrap_or_else(|instance| {
             bug!("typeid_for_instance: couldn't get fn_abi of instance {:?}", instance)
         });
 
-    // If this instance is a method and self is a reference, get the impl it belongs to
-    let impl_def_id = tcx.impl_of_method(instance.def_id());
-    if impl_def_id.is_some() && !fn_abi.args.is_empty() && fn_abi.args[0].layout.ty.is_ref() {
-        // If this impl is not an inherent impl, get the trait it implements
-        if let Some(trait_ref) = tcx.impl_trait_ref(impl_def_id.unwrap()) {
-            // Transform the concrete self into a reference to a trait object
-            let existential_predicate = trait_ref.map_bound(|trait_ref| {
-                ty::ExistentialPredicate::Trait(ty::ExistentialTraitRef::erase_self_ty(
-                    tcx, trait_ref,
-                ))
-            });
-            let existential_predicates = tcx.mk_poly_existential_predicates(&[ty::Binder::dummy(
-                existential_predicate.skip_binder(),
-            )]);
-            // Is the concrete self mutable?
-            let self_ty = if fn_abi.args[0].layout.ty.is_mutable_ptr() {
-                Ty::new_mut_ref(
-                    tcx,
-                    tcx.lifetimes.re_erased,
-                    Ty::new_dynamic(tcx, existential_predicates, tcx.lifetimes.re_erased, ty::Dyn),
-                )
-            } else {
-                Ty::new_imm_ref(
-                    tcx,
-                    tcx.lifetimes.re_erased,
-                    Ty::new_dynamic(tcx, existential_predicates, tcx.lifetimes.re_erased, ty::Dyn),
-                )
-            };
-
-            // Replace the concrete self in an fn_abi clone by the reference to a trait object
-            let mut fn_abi = fn_abi.clone();
-            // HACK(rcvalle): It is okay to not replace or update the entire ArgAbi here because the
-            //   other fields are never used.
-            fn_abi.args[0].layout.ty = self_ty;
-
-            return typeid_for_fnabi(tcx, &fn_abi, options);
-        }
-    }
-
     typeid_for_fnabi(tcx, fn_abi, options)
 }
 

tests/ui/sanitizer/cfi-complex-receiver.rs

@@ -0,0 +1,42 @@
+// Check that more complex receivers work:
+// * Arc<dyn Foo> as for custom receivers
+// * &dyn Bar<T=Baz> for type constraints
+
+//@ needs-sanitizer-cfi
+// FIXME(#122848) Remove only-linux once OSX CFI binaries work
+//@ only-linux
+//@ compile-flags: --crate-type=bin -Cprefer-dynamic=off -Clto -Zsanitizer=cfi
+//@ compile-flags: -C target-feature=-crt-static -C codegen-units=1 -C opt-level=0
+//@ run-pass
+
+use std::sync::Arc;
+
+trait Foo {
+    fn foo(self: Arc<Self>);
+}
+
+struct FooImpl;
+
+impl Foo for FooImpl {
+    fn foo(self: Arc<Self>) {}
+}
+
+trait Bar {
+    type T;
+    fn bar(&self) -> Self::T;
+}
+
+struct BarImpl;
+
+impl Bar for BarImpl {
+    type T = i32;
+    fn bar(&self) -> Self::T { 7 }
+}
+
+fn main() {
+    let foo: Arc<dyn Foo> = Arc::new(FooImpl);
+    foo.foo();
+
+    let bar: &dyn Bar<T=i32> = &BarImpl;
+    assert_eq!(bar.bar(), 7);
+}