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 async_fn_with_borrow(x: &u8) -> u8 {
  83     await!(wake_and_yield_once());
  84     *x
  85 }
  86
  87 async fn async_fn_with_borrow_named_lifetime<'a>(x: &'a u8) -> u8 {
  88     await!(wake_and_yield_once());
  89     *x
  90 }
  91
  92 fn async_fn_with_impl_future_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
  93     async move {
  94         await!(wake_and_yield_once());
  95         *x
  96     }
  97 }
  98
  99 async fn async_fn_with_named_lifetime_multiple_args<'a>(x: &'a u8, _y: &'a u8) -> u8 {
 100     await!(wake_and_yield_once());
 101     *x
 102 }
 103
 104 fn async_fn_with_internal_borrow(y: u8) -> impl Future<Output = u8> {
 105     async move {
 106         await!(async_fn_with_borrow(&y))
 107     }
 108 }
 109
 110 unsafe async fn unsafe_async_fn(x: u8) -> u8 {
 111     await!(wake_and_yield_once());
 112     x
 113 }
 114
 115 struct Foo;
 116
 117 trait Bar {
 118     fn foo() {}
 119 }
 120
 121 impl Foo {
 122     async fn async_method(x: u8) -> u8 {
 123         unsafe {
 124             await!(unsafe_async_fn(x))
 125         }
 126     }
 127 }
 128
 129 fn test_future_yields_once_then_returns<F, Fut>(f: F)
 130 where
 131     F: FnOnce(u8) -> Fut,
 132     Fut: Future<Output = u8>,
 133 {
 134     let mut fut = Box::pin(f(9));
 135     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
 136     let waker = ArcWake::into_waker(counter.clone());
 137     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
 138     assert_eq!(Poll::Pending, fut.as_mut().poll(&waker));
 139     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
 140     assert_eq!(Poll::Ready(9), fut.as_mut().poll(&waker));
 141 }
 142
 143 fn main() {
 144     macro_rules! test {
 145         ($($fn_name:expr,)*) => { $(
 146             test_future_yields_once_then_returns($fn_name);
 147         )* }
 148     }
 149
 150     macro_rules! test_with_borrow {
 151         ($($fn_name:expr,)*) => { $(
 152             test_future_yields_once_then_returns(|x| {
 153                 async move {
 154                     await!($fn_name(&x))
 155                 }
 156             });
 157         )* }
 158     }
 159
 160     test! {
 161         async_block,
 162         async_nonmove_block,
 163         async_closure,
 164         async_fn,
 165         async_fn_with_internal_borrow,
 166         Foo::async_method,
 167         |x| {
 168             async move {
 169                 unsafe { await!(unsafe_async_fn(x)) }
 170             }
 171         },
 172     }
 173
 174     test_with_borrow! {
 175         async_block_with_borrow_named_lifetime,
 176         async_fn_with_borrow,
 177         async_fn_with_borrow_named_lifetime,
 178         async_fn_with_impl_future_named_lifetime,
 179         |x| {
 180             async move {
 181                 await!(async_fn_with_named_lifetime_multiple_args(x, x))
 182             }
 183         },
 184     }
 185 }