tests/ui/async-await/async-closure.rs

   1 // run-pass
   2
   3 // revisions: default nomiropt
   4 //[nomiropt]compile-flags: -Z mir-opt-level=0
   5
   6 // edition:2018
   7 // aux-build:arc_wake.rs
   8
   9 #![feature(async_closure)]
  10
  11 extern crate arc_wake;
  12
  13 use std::pin::Pin;
  14 use std::future::Future;
  15 use std::sync::{
  16     Arc,
  17     atomic::{self, AtomicUsize},
  18 };
  19 use std::task::{Context, Poll};
  20 use arc_wake::ArcWake;
  21
  22 struct Counter {
  23     wakes: AtomicUsize,
  24 }
  25
  26 impl ArcWake for Counter {
  27     fn wake(self: Arc<Self>) {
  28         Self::wake_by_ref(&self)
  29     }
  30     fn wake_by_ref(arc_self: &Arc<Self>) {
  31         arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
  32     }
  33 }
  34
  35 struct WakeOnceThenComplete(bool);
  36
  37 fn wake_and_yield_once() -> WakeOnceThenComplete { WakeOnceThenComplete(false) }
  38
  39 impl Future for WakeOnceThenComplete {
  40     type Output = ();
  41     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
  42         if self.0 {
  43             Poll::Ready(())
  44         } else {
  45             cx.waker().wake_by_ref();
  46             self.0 = true;
  47             Poll::Pending
  48         }
  49     }
  50 }
  51
  52 fn async_closure(x: u8) -> impl Future<Output = u8> {
  53     (async move |x: u8| -> u8 {
  54         wake_and_yield_once().await;
  55         x
  56     })(x)
  57 }
  58
  59 fn async_closure_in_unsafe_block(x: u8) -> impl Future<Output = u8> {
  60     (unsafe {
  61         async move |x: u8| unsafe_fn(unsafe_async_fn(x).await)
  62     })(x)
  63 }
  64
  65 async unsafe fn unsafe_async_fn(x: u8) -> u8 {
  66     wake_and_yield_once().await;
  67     x
  68 }
  69
  70 unsafe fn unsafe_fn(x: u8) -> u8 {
  71     x
  72 }
  73
  74 fn test_future_yields_once_then_returns<F, Fut>(f: F)
  75 where
  76     F: FnOnce(u8) -> Fut,
  77     Fut: Future<Output = u8>,
  78 {
  79     let mut fut = Box::pin(f(9));
  80     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
  81     let waker = ArcWake::into_waker(counter.clone());
  82     let mut cx = Context::from_waker(&waker);
  83     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
  84     assert_eq!(Poll::Pending, fut.as_mut().poll(&mut cx));
  85     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
  86     assert_eq!(Poll::Ready(9), fut.as_mut().poll(&mut cx));
  87 }
  88
  89 fn main() {
  90     macro_rules! test {
  91         ($($fn_name:expr,)*) => { $(
  92             test_future_yields_once_then_returns($fn_name);
  93         )* }
  94     }
  95
  96     test! {
  97         async_closure,
  98         async_closure_in_unsafe_block,
  99     }
 100 }