CFI: Rewrite closure and coroutine instances to their trait method

Similar to methods on a trait object, the most common way to indirectly
call a closure or coroutine is through the vtable on the appropriate
trait. This uses the same approach as we use for trait methods, after
backing out the trait arguments from the type.

Author: maurer

Diff for: compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs

@@ -10,6 +10,7 @@
 use rustc_data_structures::base_n;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir as hir;
+use rustc_hir::lang_items::LangItem;
 use rustc_middle::ty::layout::IntegerExt;
 use rustc_middle::ty::TypeVisitableExt;
 use rustc_middle::ty::{
@@ -1152,6 +1153,14 @@ pub fn typeid_for_instance<'tcx>(
         instance.args = tcx.mk_args_trait(self_ty, List::empty());
     } else if matches!(instance.def, ty::InstanceDef::Virtual(..)) {
         instance.args = strip_receiver_auto(tcx, instance.args);
+    } else if let ty::InstanceDef::VTableShim(def_id) = instance.def
+        && let Some(trait_id) = tcx.trait_of_item(def_id)
+    {
+        // VTableShims may have a trait method, but a concrete Self. This is not suitable for a vtable,
+        // as the caller will not know the concrete Self.
+        let trait_ref = ty::TraitRef::new(tcx, trait_id, instance.args);
+        let invoke_ty = trait_object_ty(tcx, ty::Binder::dummy(trait_ref));
+        instance.args = tcx.mk_args_trait(invoke_ty, trait_ref.args.into_iter().skip(1));
     }
 
     if !options.contains(EncodeTyOptions::NO_SELF_TYPE_ERASURE)
@@ -1190,6 +1199,45 @@ pub fn typeid_for_instance<'tcx>(
                 tcx.mk_args_trait(invoke_ty, trait_ref.args.into_iter().skip(1));
             instance.args = instance.args.rebase_onto(tcx, impl_id, abstract_trait_args);
         }
+    } else if tcx.is_closure_like(instance.def_id()) {
+        // We're either a closure or a coroutine. Our goal is to find the trait we're defined on,
+        // instantiate it, and take the type of its only method as our own.
+        let closure_ty = instance.ty(tcx, ty::ParamEnv::reveal_all());
+        let (trait_id, inputs) = match closure_ty.kind() {
+            ty::Closure(..) => {
+                let closure_args = instance.args.as_closure();
+                let trait_id = tcx.fn_trait_kind_to_def_id(closure_args.kind()).unwrap();
+                let tuple_args =
+                    tcx.instantiate_bound_regions_with_erased(closure_args.sig()).inputs()[0];
+                (trait_id, tuple_args)
+            }
+            ty::Coroutine(..) => (
+                tcx.require_lang_item(LangItem::Coroutine, None),
+                instance.args.as_coroutine().resume_ty(),
+            ),
+            ty::CoroutineClosure(..) => (
+                tcx.require_lang_item(LangItem::FnOnce, None),
+                tcx.instantiate_bound_regions_with_erased(
+                    instance.args.as_coroutine_closure().coroutine_closure_sig(),
+                )
+                .tupled_inputs_ty,
+            ),
+            x => bug!("Unexpected type kind for closure-like: {x:?}"),
+        };
+        let trait_ref = ty::TraitRef::new(tcx, trait_id, [closure_ty, inputs]);
+        let invoke_ty = trait_object_ty(tcx, ty::Binder::dummy(trait_ref));
+        let abstract_args = tcx.mk_args_trait(invoke_ty, trait_ref.args.into_iter().skip(1));
+        // There should be exactly one method on this trait, and it should be the one we're
+        // defining.
+        let call = tcx
+            .associated_items(trait_id)
+            .in_definition_order()
+            .find(|it| it.kind == ty::AssocKind::Fn)
+            .expect("No call-family function on closure-like Fn trait?")
+            .def_id;
+
+        instance.def = ty::InstanceDef::Virtual(call, 0);
+        instance.args = abstract_args;
     }
 
     let fn_abi = tcx

Diff for: tests/ui/sanitizer/cfi-async-closures.rs

@@ -10,6 +10,7 @@
 //@ [cfi] compile-flags: -C codegen-units=1 -C lto -C prefer-dynamic=off -C opt-level=0
 //@ [cfi] compile-flags: -Z sanitizer=cfi
 //@ [kcfi] compile-flags: -Z sanitizer=kcfi
+//@ [kcfi] compile-flags: -C panic=abort -Z panic-abort-tests -C prefer-dynamic=off
 //@ run-pass
 
 #![feature(async_closure)]
@@ -27,4 +28,6 @@ fn main() {
    let f = identity(async || ());
    let _ = f.async_call(());
    let _ = f();
+   let g: Box<dyn FnOnce() -> _> = Box::new(f) as _;
+   let _ = g();
 }

Diff for: tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs

@@ -0,0 +1,66 @@
+// Check various forms of dynamic closure calls
+
+//@ revisions: cfi kcfi
+// FIXME(#122848) Remove only-linux once OSX CFI binaries work
+//@ only-linux
+//@ [cfi] needs-sanitizer-cfi
+//@ [kcfi] needs-sanitizer-kcfi
+//@ compile-flags: -C target-feature=-crt-static
+//@ [cfi] compile-flags: -C codegen-units=1 -C lto -C prefer-dynamic=off -C opt-level=0
+//@ [cfi] compile-flags: -Z sanitizer=cfi
+//@ [kcfi] compile-flags: -Z sanitizer=kcfi
+//@ [kcfi] compile-flags: -C panic=abort -Z panic-abort-tests -C prefer-dynamic=off
+//@ compile-flags: --test
+//@ run-pass
+
+#![feature(fn_traits)]
+
+fn foo<'a, T>() -> Box<dyn Fn(&'a T) -> &'a T> {
+    Box::new(|x| x)
+}
+
+#[test]
+fn dyn_fn_with_params() {
+    let x = 3;
+    let f = foo();
+    f(&x);
+    // FIXME remove once drops are working.
+    std::mem::forget(f);
+}
+
+#[test]
+fn call_fn_trait() {
+   let f: &(dyn Fn()) = &(|| {}) as _;
+   f.call(());
+}
+
+#[test]
+fn fn_ptr_cast() {
+    let f: &fn() = &((|| ()) as _);
+    f();
+}
+
+fn use_fnmut<F: FnMut()>(mut f: F) {
+    f()
+}
+
+#[test]
+fn fn_to_fnmut() {
+    let f: &(dyn Fn()) = &(|| {}) as _;
+    use_fnmut(f);
+}
+
+fn hrtb_helper(f: &dyn for<'a> Fn(&'a usize)) {
+    f(&10)
+}
+
+#[test]
+fn hrtb_fn() {
+    hrtb_helper((&|x: &usize| println!("{}", *x)) as _)
+}
+
+#[test]
+fn fnonce() {
+    let f: Box<dyn FnOnce()> = Box::new(|| {}) as _;
+    f();
+}

Diff for: tests/ui/sanitizer/cfi-closures.rs

@@ -0,0 +1,29 @@
+// Verifies that we can call dynamic coroutines
+
+//@ revisions: cfi kcfi
+// FIXME(#122848) Remove only-linux once OSX CFI binaries work
+//@ only-linux
+//@ [cfi] needs-sanitizer-cfi
+//@ [kcfi] needs-sanitizer-kcfi
+//@ compile-flags: -C target-feature=-crt-static
+//@ [cfi] compile-flags: -C codegen-units=1 -C lto -C prefer-dynamic=off -C opt-level=0
+//@ [cfi] compile-flags: -Z sanitizer=cfi
+//@ [kcfi] compile-flags: -Z sanitizer=kcfi
+//@ compile-flags: --test
+//@ run-pass
+
+#![feature(coroutines)]
+#![feature(coroutine_trait)]
+
+use std::ops::{Coroutine, CoroutineState};
+use std::pin::{pin, Pin};
+
+fn main() {
+    let mut coro = |x: i32| {
+        yield x;
+        "done"
+    };
+    let mut abstract_coro: Pin<&mut dyn Coroutine<i32,Yield=i32,Return=&'static str>> = pin!(coro);
+    assert_eq!(abstract_coro.as_mut().resume(2), CoroutineState::Yielded(2));
+    assert_eq!(abstract_coro.as_mut().resume(0), CoroutineState::Complete("done"));
+}