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

   1 // run-pass
   2
   3 // edition:2018
   4 // aux-build:arc_wake.rs
   5
   6 #![feature(async_await)]
   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_block(x: u8) -> impl Future<Output = u8> {
  50     async move {
  51         wake_and_yield_once().await;
  52         x
  53     }
  54 }
  55
  56 fn async_block_with_borrow_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
  57     async move {
  58         wake_and_yield_once().await;
  59         *x
  60     }
  61 }
  62
  63 fn async_nonmove_block(x: u8) -> impl Future<Output = u8> {
  64     async move {
  65         let future = async {
  66             wake_and_yield_once().await;
  67             x
  68         };
  69         future.await
  70     }
  71 }
  72
  73 fn async_closure(x: u8) -> impl Future<Output = u8> {
  74     (async move |x: u8| -> u8 {
  75         wake_and_yield_once().await;
  76         x
  77     })(x)
  78 }
  79
  80 async fn async_fn(x: u8) -> u8 {
  81     wake_and_yield_once().await;
  82     x
  83 }
  84
  85 async fn generic_async_fn<T>(x: T) -> T {
  86     wake_and_yield_once().await;
  87     x
  88 }
  89
  90 async fn async_fn_with_borrow(x: &u8) -> u8 {
  91     wake_and_yield_once().await;
  92     *x
  93 }
  94
  95 async fn async_fn_with_borrow_named_lifetime<'a>(x: &'a u8) -> u8 {
  96     wake_and_yield_once().await;
  97     *x
  98 }
  99
 100 fn async_fn_with_impl_future_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
 101     async move {
 102         wake_and_yield_once().await;
 103         *x
 104     }
 105 }
 106
 107 /* FIXME(cramertj) support when `existential type T<'a, 'b>:;` works
 108 async fn async_fn_multiple_args(x: &u8, _y: &u8) -> u8 {
 109     await!(wake_and_yield_once());
 110     *x
 111 }
 112 */
 113
 114 async fn async_fn_multiple_args_named_lifetime<'a>(x: &'a u8, _y: &'a u8) -> u8 {
 115     wake_and_yield_once().await;
 116     *x
 117 }
 118
 119 fn async_fn_with_internal_borrow(y: u8) -> impl Future<Output = u8> {
 120     async move {
 121         async_fn_with_borrow_named_lifetime(&y).await
 122     }
 123 }
 124
 125 async unsafe fn unsafe_async_fn(x: u8) -> u8 {
 126     wake_and_yield_once().await;
 127     x
 128 }
 129
 130 struct Foo;
 131
 132 trait Bar {
 133     fn foo() {}
 134 }
 135
 136 impl Foo {
 137     async fn async_assoc_item(x: u8) -> u8 {
 138         unsafe {
 139             unsafe_async_fn(x).await
 140         }
 141     }
 142
 143     async unsafe fn async_unsafe_assoc_item(x: u8) -> u8 {
 144         unsafe_async_fn(x).await
 145     }
 146 }
 147
 148 fn test_future_yields_once_then_returns<F, Fut>(f: F)
 149 where
 150     F: FnOnce(u8) -> Fut,
 151     Fut: Future<Output = u8>,
 152 {
 153     let mut fut = Box::pin(f(9));
 154     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
 155     let waker = ArcWake::into_waker(counter.clone());
 156     let mut cx = Context::from_waker(&waker);
 157     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
 158     assert_eq!(Poll::Pending, fut.as_mut().poll(&mut cx));
 159     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
 160     assert_eq!(Poll::Ready(9), fut.as_mut().poll(&mut cx));
 161 }
 162
 163 fn main() {
 164     macro_rules! test {
 165         ($($fn_name:expr,)*) => { $(
 166             test_future_yields_once_then_returns($fn_name);
 167         )* }
 168     }
 169
 170     macro_rules! test_with_borrow {
 171         ($($fn_name:expr,)*) => { $(
 172             test_future_yields_once_then_returns(|x| {
 173                 async move {
 174                     $fn_name(&x).await
 175                 }
 176             });
 177         )* }
 178     }
 179
 180     test! {
 181         async_block,
 182         async_nonmove_block,
 183         async_closure,
 184         async_fn,
 185         generic_async_fn,
 186         async_fn_with_internal_borrow,
 187         Foo::async_assoc_item,
 188         |x| {
 189             async move {
 190                 unsafe { unsafe_async_fn(x).await }
 191             }
 192         },
 193         |x| {
 194             async move {
 195                 unsafe { Foo::async_unsafe_assoc_item(x).await }
 196             }
 197         },
 198     }
 199     test_with_borrow! {
 200         async_block_with_borrow_named_lifetime,
 201         async_fn_with_borrow,
 202         async_fn_with_borrow_named_lifetime,
 203         async_fn_with_impl_future_named_lifetime,
 204         |x| {
 205             async move {
 206                 async_fn_multiple_args_named_lifetime(x, x).await
 207             }
 208         },
 209     }
 210 }