CFI: Support non-general coroutines

Previously, we assumed all `ty::Coroutine` were general coroutines and
attempted to generalize them through the `Coroutine` trait. Select
appropriate traits for each kind of coroutine.
This commit is contained in:
Matthew Maurer 2024-04-02 15:17:33 +00:00
parent 6bb6b816bf
commit a333b82d04
2 changed files with 62 additions and 12 deletions

View File

@ -1218,22 +1218,35 @@ pub fn typeid_for_instance<'tcx>(
let trait_id = tcx.fn_trait_kind_to_def_id(closure_args.kind()).unwrap(); let trait_id = tcx.fn_trait_kind_to_def_id(closure_args.kind()).unwrap();
let tuple_args = let tuple_args =
tcx.instantiate_bound_regions_with_erased(closure_args.sig()).inputs()[0]; tcx.instantiate_bound_regions_with_erased(closure_args.sig()).inputs()[0];
(trait_id, tuple_args) (trait_id, Some(tuple_args))
} }
ty::Coroutine(..) => ( ty::Coroutine(..) => match tcx.coroutine_kind(instance.def_id()).unwrap() {
tcx.require_lang_item(LangItem::Coroutine, None), hir::CoroutineKind::Coroutine(..) => (
instance.args.as_coroutine().resume_ty(), tcx.require_lang_item(LangItem::Coroutine, None),
), Some(instance.args.as_coroutine().resume_ty()),
),
hir::CoroutineKind::Desugared(desugaring, _) => {
let lang_item = match desugaring {
hir::CoroutineDesugaring::Async => LangItem::Future,
hir::CoroutineDesugaring::AsyncGen => LangItem::AsyncIterator,
hir::CoroutineDesugaring::Gen => LangItem::Iterator,
};
(tcx.require_lang_item(lang_item, None), None)
}
},
ty::CoroutineClosure(..) => ( ty::CoroutineClosure(..) => (
tcx.require_lang_item(LangItem::FnOnce, None), tcx.require_lang_item(LangItem::FnOnce, None),
tcx.instantiate_bound_regions_with_erased( Some(
instance.args.as_coroutine_closure().coroutine_closure_sig(), tcx.instantiate_bound_regions_with_erased(
) instance.args.as_coroutine_closure().coroutine_closure_sig(),
.tupled_inputs_ty, )
.tupled_inputs_ty,
),
), ),
x => bug!("Unexpected type kind for closure-like: {x:?}"), x => bug!("Unexpected type kind for closure-like: {x:?}"),
}; };
let trait_ref = ty::TraitRef::new(tcx, trait_id, [closure_ty, inputs]); let concrete_args = tcx.mk_args_trait(closure_ty, inputs.map(Into::into));
let trait_ref = ty::TraitRef::new(tcx, trait_id, concrete_args);
let invoke_ty = trait_object_ty(tcx, ty::Binder::dummy(trait_ref)); 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)); 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 // There should be exactly one method on this trait, and it should be the one we're

View File

@ -3,6 +3,7 @@
//@ revisions: cfi kcfi //@ revisions: cfi kcfi
// FIXME(#122848) Remove only-linux once OSX CFI binaries work // FIXME(#122848) Remove only-linux once OSX CFI binaries work
//@ only-linux //@ only-linux
//@ edition: 2024
//@ [cfi] needs-sanitizer-cfi //@ [cfi] needs-sanitizer-cfi
//@ [kcfi] needs-sanitizer-kcfi //@ [kcfi] needs-sanitizer-kcfi
//@ compile-flags: -C target-feature=-crt-static //@ compile-flags: -C target-feature=-crt-static
@ -10,16 +11,22 @@
//@ [cfi] compile-flags: -Z sanitizer=cfi //@ [cfi] compile-flags: -Z sanitizer=cfi
//@ [kcfi] compile-flags: -Z sanitizer=kcfi //@ [kcfi] compile-flags: -Z sanitizer=kcfi
//@ [kcfi] compile-flags: -C panic=abort -Z panic-abort-tests -C prefer-dynamic=off //@ [kcfi] compile-flags: -C panic=abort -Z panic-abort-tests -C prefer-dynamic=off
//@ compile-flags: --test //@ compile-flags: --test -Z unstable-options
//@ run-pass //@ run-pass
#![feature(coroutines)] #![feature(coroutines)]
#![feature(coroutine_trait)] #![feature(coroutine_trait)]
#![feature(noop_waker)]
#![feature(gen_blocks)]
#![feature(async_iterator)]
use std::ops::{Coroutine, CoroutineState}; use std::ops::{Coroutine, CoroutineState};
use std::pin::{pin, Pin}; use std::pin::{pin, Pin};
use std::task::{Context, Poll, Waker};
use std::async_iter::AsyncIterator;
fn main() { #[test]
fn general_coroutine() {
let mut coro = |x: i32| { let mut coro = |x: i32| {
yield x; yield x;
"done" "done"
@ -28,3 +35,33 @@ fn main() {
assert_eq!(abstract_coro.as_mut().resume(2), CoroutineState::Yielded(2)); assert_eq!(abstract_coro.as_mut().resume(2), CoroutineState::Yielded(2));
assert_eq!(abstract_coro.as_mut().resume(0), CoroutineState::Complete("done")); assert_eq!(abstract_coro.as_mut().resume(0), CoroutineState::Complete("done"));
} }
async fn async_fn() {}
#[test]
fn async_coroutine() {
let f: fn() -> Pin<Box<dyn Future<Output = ()>>> = || Box::pin(async_fn());
let _ = async { f().await; };
assert_eq!(f().as_mut().poll(&mut Context::from_waker(Waker::noop())), Poll::Ready(()));
}
async gen fn async_gen_fn() -> u8 {
yield 5;
}
#[test]
fn async_gen_coroutine() {
let f: fn() -> Pin<Box<dyn AsyncIterator<Item = u8>>> = || Box::pin(async_gen_fn());
assert_eq!(f().as_mut().poll_next(&mut Context::from_waker(Waker::noop())),
Poll::Ready(Some(5)));
}
gen fn gen_fn() -> u8 {
yield 6;
}
#[test]
fn gen_coroutine() {
let f: fn() -> Box<dyn Iterator<Item = u8>> = || Box::new(gen_fn());
assert_eq!(f().next(), Some(6));
}