src/test/ui/async-await/async-closure.rs

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