Auto merge of #123106 - maurer:cfi-closures, r=compiler-errors

CFI: Abstract Closures and Coroutines

This will abstract coroutines in a moment, it's just abstracting closures for now to show `@rcvalle`

This uses the same principal as the methods on traits - figure out the `dyn` type representing the fn trait, instantiate it, and attach that alias set. We're essentially just computing how we would be called in a dynamic context, and attaching that.

5 files changed, 240 insertions, 58 deletions

diff --git a/compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs b/compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs
index aba3212b4ce..5963bd7c5f1 100644
--- a/compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs
+++ b/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::{
@@ -641,9 +642,7 @@ fn encode_ty<'tcx>(
         }
 
         // Function types
-        ty::FnDef(def_id, args)
-        | ty::Closure(def_id, args)
-        | ty::CoroutineClosure(def_id, args) => {
+        ty::FnDef(def_id, args) | ty::Closure(def_id, args) => {
             // u<length><name>[I<element-type1..element-typeN>E], where <element-type> is <subst>,
             // as vendor extended type.
             let mut s = String::new();
@@ -654,6 +653,18 @@ fn encode_ty<'tcx>(
             typeid.push_str(&s);
         }
 
+        ty::CoroutineClosure(def_id, args) => {
+            // u<length><name>[I<element-type1..element-typeN>E], where <element-type> is <subst>,
+            // as vendor extended type.
+            let mut s = String::new();
+            let name = encode_ty_name(tcx, *def_id);
+            let _ = write!(s, "u{}{}", name.len(), &name);
+            let parent_args = tcx.mk_args(args.as_coroutine_closure().parent_args());
+            s.push_str(&encode_args(tcx, parent_args, dict, options));
+            compress(dict, DictKey::Ty(ty, TyQ::None), &mut s);
+            typeid.push_str(&s);
+        }
+
         ty::Coroutine(def_id, args, ..) => {
             // u<length><name>[I<element-type1..element-typeN>E], where <element-type> is <subst>,
             // as vendor extended type.
@@ -1151,43 +1162,91 @@ pub fn typeid_for_instance<'tcx>(
         };
         let stripped_ty = strip_receiver_auto(tcx, upcast_ty);
         instance.args = tcx.mk_args_trait(stripped_ty, instance.args.into_iter().skip(1));
+    } 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)
-        && let Some(impl_id) = tcx.impl_of_method(instance.def_id())
-        && let Some(trait_ref) = tcx.impl_trait_ref(impl_id)
-    {
-        let impl_method = tcx.associated_item(instance.def_id());
-        let method_id = impl_method
-            .trait_item_def_id
-            .expect("Part of a trait implementation, but not linked to the def_id?");
-        let trait_method = tcx.associated_item(method_id);
-        let trait_id = trait_ref.skip_binder().def_id;
-        if traits::is_vtable_safe_method(tcx, trait_id, trait_method)
-            && tcx.object_safety_violations(trait_id).is_empty()
+    if !options.contains(EncodeTyOptions::NO_SELF_TYPE_ERASURE) {
+        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 = tcx.instantiate_and_normalize_erasing_regions(
-                instance.args,
-                ty::ParamEnv::reveal_all(),
-                trait_ref,
-            );
+            let impl_method = tcx.associated_item(instance.def_id());
+            let method_id = impl_method
+                .trait_item_def_id
+                .expect("Part of a trait implementation, but not linked to the def_id?");
+            let trait_method = tcx.associated_item(method_id);
+            let trait_id = trait_ref.skip_binder().def_id;
+            if traits::is_vtable_safe_method(tcx, trait_id, trait_method)
+                && tcx.object_safety_violations(trait_id).is_empty()
+            {
+                // 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 = tcx.instantiate_and_normalize_erasing_regions(
+                    instance.args,
+                    ty::ParamEnv::reveal_all(),
+                    trait_ref,
+                );
+                let invoke_ty = trait_object_ty(tcx, ty::Binder::dummy(trait_ref));
+
+                // At the call site, any call to this concrete function through a vtable will be
+                // `Virtual(method_id, idx)` with appropriate arguments for the method. Since we have the
+                // original method id, and we've recovered the trait arguments, we can make the callee
+                // instance we're computing the alias set for match the caller instance.
+                //
+                // Right now, our code ignores the vtable index everywhere, so we use 0 as a placeholder.
+                // If we ever *do* start encoding the vtable index, we will need to generate an alias set
+                // based on which vtables we are putting this method into, as there will be more than one
+                // index value when supertraits are involved.
+                instance.def = ty::InstanceDef::Virtual(method_id, 0);
+                let abstract_trait_args =
+                    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;
 
-            // At the call site, any call to this concrete function through a vtable will be
-            // `Virtual(method_id, idx)` with appropriate arguments for the method. Since we have the
-            // original method id, and we've recovered the trait arguments, we can make the callee
-            // instance we're computing the alias set for match the caller instance.
-            //
-            // Right now, our code ignores the vtable index everywhere, so we use 0 as a placeholder.
-            // If we ever *do* start encoding the vtable index, we will need to generate an alias set
-            // based on which vtables we are putting this method into, as there will be more than one
-            // index value when supertraits are involved.
-            instance.def = ty::InstanceDef::Virtual(method_id, 0);
-            let abstract_trait_args =
-                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);
+            instance.def = ty::InstanceDef::Virtual(call, 0);
+            instance.args = abstract_args;
         }
     }
 
diff --git a/tests/ui/sanitizer/cfi-async-closures.rs b/tests/ui/sanitizer/cfi-async-closures.rs
new file mode 100644
index 00000000000..d94f2992d84
--- /dev/null
+++ b/tests/ui/sanitizer/cfi-async-closures.rs
@@ -0,0 +1,33 @@
+// Check various forms of dynamic closure calls
+
+//@ edition: 2021
+//@ 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
+//@ run-pass
+
+#![feature(async_closure)]
+#![feature(async_fn_traits)]
+
+use std::ops::AsyncFn;
+
+#[inline(never)]
+fn identity<T>(x: T) -> T { x }
+
+// We can't actually create a `dyn AsyncFn()`, because it's not object-safe, but we should check
+// that we don't bug out when we encounter one.
+
+fn main() {
+   let f = identity(async || ());
+   let _ = f.async_call(());
+   let _ = f();
+   let g: Box<dyn FnOnce() -> _> = Box::new(f) as _;
+   let _ = g();
+}
diff --git a/tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs b/tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs
deleted file mode 100644
index 03818544aab..00000000000
--- a/tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs
+++ /dev/null
@@ -1,23 +0,0 @@
-// Tests that converting a closure to a function pointer works
-// The notable thing being tested here is that when the closure does not capture anything,
-// the call method from its Fn trait takes a ZST representing its environment. The compiler then
-// uses the assumption that the ZST is non-passed to reify this into a function pointer.
-//
-// This checks that the reified function pointer will have the expected alias set at its call-site.
-
-//@ 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 -C prefer-dynamic=off
-//@ run-pass
-
-pub fn main() {
-    let f: &fn() = &((|| ()) as _);
-    f();
-}
diff --git a/tests/ui/sanitizer/cfi-closures.rs b/tests/ui/sanitizer/cfi-closures.rs
new file mode 100644
index 00000000000..54f1cc035bc
--- /dev/null
+++ b/tests/ui/sanitizer/cfi-closures.rs
@@ -0,0 +1,83 @@
+// 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)]
+#![feature(unboxed_closures)]
+#![feature(cfg_sanitize)]
+
+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();
+}
+
+fn use_closure<C>(call: extern "rust-call" fn(&C, ()) -> i32, f: &C) -> i32 {
+    call(f, ())
+}
+
+#[test]
+// FIXME after KCFI reify support is added, remove this
+// It will appear to work if you test locally, set -C opt-level=0 to see it fail.
+#[cfg_attr(sanitize = "kcfi", ignore)]
+fn closure_addr_taken() {
+    let x = 3i32;
+    let f = || x;
+    let call = Fn::<()>::call;
+    use_closure(call, &f);
+}
diff --git a/tests/ui/sanitizer/cfi-coroutine.rs b/tests/ui/sanitizer/cfi-coroutine.rs
new file mode 100644
index 00000000000..24e59cf5b4d
--- /dev/null
+++ b/tests/ui/sanitizer/cfi-coroutine.rs
@@ -0,0 +1,30 @@
+// 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
+//@ [kcfi] compile-flags: -C panic=abort -Z panic-abort-tests -C prefer-dynamic=off
+//@ 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"));
+}