101 lines
2.2 KiB
Rust
101 lines
2.2 KiB
Rust
// run-pass
|
|
|
|
// revisions: default nomiropt
|
|
//[nomiropt]compile-flags: -Z mir-opt-level=0
|
|
|
|
// edition:2018
|
|
// aux-build:arc_wake.rs
|
|
|
|
#![feature(async_closure)]
|
|
|
|
extern crate arc_wake;
|
|
|
|
use std::pin::Pin;
|
|
use std::future::Future;
|
|
use std::sync::{
|
|
Arc,
|
|
atomic::{self, AtomicUsize},
|
|
};
|
|
use std::task::{Context, Poll};
|
|
use arc_wake::ArcWake;
|
|
|
|
struct Counter {
|
|
wakes: AtomicUsize,
|
|
}
|
|
|
|
impl ArcWake for Counter {
|
|
fn wake(self: Arc<Self>) {
|
|
Self::wake_by_ref(&self)
|
|
}
|
|
fn wake_by_ref(arc_self: &Arc<Self>) {
|
|
arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
|
|
}
|
|
}
|
|
|
|
struct WakeOnceThenComplete(bool);
|
|
|
|
fn wake_and_yield_once() -> WakeOnceThenComplete { WakeOnceThenComplete(false) }
|
|
|
|
impl Future for WakeOnceThenComplete {
|
|
type Output = ();
|
|
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
|
|
if self.0 {
|
|
Poll::Ready(())
|
|
} else {
|
|
cx.waker().wake_by_ref();
|
|
self.0 = true;
|
|
Poll::Pending
|
|
}
|
|
}
|
|
}
|
|
|
|
fn async_closure(x: u8) -> impl Future<Output = u8> {
|
|
(async move |x: u8| -> u8 {
|
|
wake_and_yield_once().await;
|
|
x
|
|
})(x)
|
|
}
|
|
|
|
fn async_closure_in_unsafe_block(x: u8) -> impl Future<Output = u8> {
|
|
(unsafe {
|
|
async move |x: u8| unsafe_fn(unsafe_async_fn(x).await)
|
|
})(x)
|
|
}
|
|
|
|
async unsafe fn unsafe_async_fn(x: u8) -> u8 {
|
|
wake_and_yield_once().await;
|
|
x
|
|
}
|
|
|
|
unsafe fn unsafe_fn(x: u8) -> u8 {
|
|
x
|
|
}
|
|
|
|
fn test_future_yields_once_then_returns<F, Fut>(f: F)
|
|
where
|
|
F: FnOnce(u8) -> Fut,
|
|
Fut: Future<Output = u8>,
|
|
{
|
|
let mut fut = Box::pin(f(9));
|
|
let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
|
|
let waker = ArcWake::into_waker(counter.clone());
|
|
let mut cx = Context::from_waker(&waker);
|
|
assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
|
|
assert_eq!(Poll::Pending, fut.as_mut().poll(&mut cx));
|
|
assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
|
|
assert_eq!(Poll::Ready(9), fut.as_mut().poll(&mut cx));
|
|
}
|
|
|
|
fn main() {
|
|
macro_rules! test {
|
|
($($fn_name:expr,)*) => { $(
|
|
test_future_yields_once_then_returns($fn_name);
|
|
)* }
|
|
}
|
|
|
|
test! {
|
|
async_closure,
|
|
async_closure_in_unsafe_block,
|
|
}
|
|
}
|