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

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