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.

4 files changed, 116 insertions, 23 deletions

diff --git a/tests/ui/sanitizer/cfi-async-closures.rs b/tests/ui/sanitizer/cfi-async-closures.rs
index 50edeb852bd..d94f2992d84 100644
--- a/tests/ui/sanitizer/cfi-async-closures.rs
+++ b/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 --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"));
+}