src/test/run-pass/async-await.rs

   1 // edition:2018
   2
   3 #![feature(arbitrary_self_types, async_await, await_macro, futures_api)]
   4
   5 use std::pin::Pin;
   6 use std::future::Future;
   7 use std::sync::{
   8     Arc,
   9     atomic::{self, AtomicUsize},
  10 };
  11 use std::task::{
  12     LocalWaker, Poll, Wake,
  13     local_waker_from_nonlocal,
  14 };
  15
  16 struct Counter {
  17     wakes: AtomicUsize,
  18 }
  19
  20 impl Wake for Counter {
  21     fn wake(this: &Arc<Self>) {
  22         this.wakes.fetch_add(1, atomic::Ordering::SeqCst);
  23     }
  24 }
  25
  26 struct WakeOnceThenComplete(bool);
  27
  28 fn wake_and_yield_once() -> WakeOnceThenComplete { WakeOnceThenComplete(false) }
  29
  30 impl Future for WakeOnceThenComplete {
  31     type Output = ();
  32     fn poll(mut self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<()> {
  33         if self.0 {
  34             Poll::Ready(())
  35         } else {
  36             lw.wake();
  37             self.0 = true;
  38             Poll::Pending
  39         }
  40     }
  41 }
  42
  43 fn async_block(x: u8) -> impl Future<Output = u8> {
  44     async move {
  45         await!(wake_and_yield_once());
  46         x
  47     }
  48 }
  49
  50 fn async_block_with_borrow_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
  51     async move {
  52         await!(wake_and_yield_once());
  53         *x
  54     }
  55 }
  56
  57 fn async_nonmove_block(x: u8) -> impl Future<Output = u8> {
  58     async move {
  59         let future = async {
  60             await!(wake_and_yield_once());
  61             x
  62         };
  63         await!(future)
  64     }
  65 }
  66
  67 fn async_closure(x: u8) -> impl Future<Output = u8> {
  68     (async move |x: u8| -> u8 {
  69         await!(wake_and_yield_once());
  70         x
  71     })(x)
  72 }
  73
  74 async fn async_fn(x: u8) -> u8 {
  75     await!(wake_and_yield_once());
  76     x
  77 }
  78
  79 async fn async_fn_with_borrow(x: &u8) -> u8 {
  80     await!(wake_and_yield_once());
  81     *x
  82 }
  83
  84 async fn async_fn_with_borrow_named_lifetime<'a>(x: &'a u8) -> u8 {
  85     await!(wake_and_yield_once());
  86     *x
  87 }
  88
  89 fn async_fn_with_impl_future_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
  90     async move {
  91         await!(wake_and_yield_once());
  92         *x
  93     }
  94 }
  95
  96 async fn async_fn_with_named_lifetime_multiple_args<'a>(x: &'a u8, _y: &'a u8) -> u8 {
  97     await!(wake_and_yield_once());
  98     *x
  99 }
 100
 101 fn async_fn_with_internal_borrow(y: u8) -> impl Future<Output = u8> {
 102     async move {
 103         await!(async_fn_with_borrow(&y))
 104     }
 105 }
 106
 107 unsafe async fn unsafe_async_fn(x: u8) -> u8 {
 108     await!(wake_and_yield_once());
 109     x
 110 }
 111
 112 struct Foo;
 113
 114 trait Bar {
 115     fn foo() {}
 116 }
 117
 118 impl Foo {
 119     async fn async_method(x: u8) -> u8 {
 120         unsafe {
 121             await!(unsafe_async_fn(x))
 122         }
 123     }
 124 }
 125
 126 fn test_future_yields_once_then_returns<F, Fut>(f: F)
 127 where
 128     F: FnOnce(u8) -> Fut,
 129     Fut: Future<Output = u8>,
 130 {
 131     let mut fut = Box::pin(f(9));
 132     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
 133     let waker = local_waker_from_nonlocal(counter.clone());
 134     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
 135     assert_eq!(Poll::Pending, fut.as_mut().poll(&waker));
 136     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
 137     assert_eq!(Poll::Ready(9), fut.as_mut().poll(&waker));
 138 }
 139
 140 fn main() {
 141     macro_rules! test {
 142         ($($fn_name:expr,)*) => { $(
 143             test_future_yields_once_then_returns($fn_name);
 144         )* }
 145     }
 146
 147     macro_rules! test_with_borrow {
 148         ($($fn_name:expr,)*) => { $(
 149             test_future_yields_once_then_returns(|x| {
 150                 async move {
 151                     await!($fn_name(&x))
 152                 }
 153             });
 154         )* }
 155     }
 156
 157     test! {
 158         async_block,
 159         async_nonmove_block,
 160         async_closure,
 161         async_fn,
 162         async_fn_with_internal_borrow,
 163         Foo::async_method,
 164         |x| {
 165             async move {
 166                 unsafe { await!(unsafe_async_fn(x)) }
 167             }
 168         },
 169     }
 170
 171     test_with_borrow! {
 172         async_block_with_borrow_named_lifetime,
 173         async_fn_with_borrow,
 174         async_fn_with_borrow_named_lifetime,
 175         async_fn_with_impl_future_named_lifetime,
 176         |x| {
 177             async move {
 178                 await!(async_fn_with_named_lifetime_multiple_args(x, x))
 179             }
 180         },
 181     }
 182 }