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

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