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 async fn async_fn(x: u8) -> u8 {
  74     wake_and_yield_once().await;
  75     x
  76 }
  77
  78 async fn generic_async_fn<T>(x: T) -> T {
  79     wake_and_yield_once().await;
  80     x
  81 }
  82
  83 async fn async_fn_with_borrow(x: &u8) -> u8 {
  84     wake_and_yield_once().await;
  85     *x
  86 }
  87
  88 async fn async_fn_with_borrow_named_lifetime<'a>(x: &'a u8) -> u8 {
  89     wake_and_yield_once().await;
  90     *x
  91 }
  92
  93 fn async_fn_with_impl_future_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
  94     async move {
  95         wake_and_yield_once().await;
  96         *x
  97     }
  98 }
  99
 100 /* FIXME(cramertj) support when `existential type T<'a, 'b>:;` works
 101 async fn async_fn_multiple_args(x: &u8, _y: &u8) -> u8 {
 102     await!(wake_and_yield_once());
 103     *x
 104 }
 105 */
 106
 107 async fn async_fn_multiple_args_named_lifetime<'a>(x: &'a u8, _y: &'a u8) -> u8 {
 108     wake_and_yield_once().await;
 109     *x
 110 }
 111
 112 fn async_fn_with_internal_borrow(y: u8) -> impl Future<Output = u8> {
 113     async move {
 114         async_fn_with_borrow_named_lifetime(&y).await
 115     }
 116 }
 117
 118 async unsafe fn unsafe_async_fn(x: u8) -> u8 {
 119     wake_and_yield_once().await;
 120     x
 121 }
 122
 123 struct Foo;
 124
 125 trait Bar {
 126     fn foo() {}
 127 }
 128
 129 impl Foo {
 130     async fn async_assoc_item(x: u8) -> u8 {
 131         unsafe {
 132             unsafe_async_fn(x).await
 133         }
 134     }
 135
 136     async unsafe fn async_unsafe_assoc_item(x: u8) -> u8 {
 137         unsafe_async_fn(x).await
 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     let mut cx = Context::from_waker(&waker);
 150     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
 151     assert_eq!(Poll::Pending, fut.as_mut().poll(&mut cx));
 152     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
 153     assert_eq!(Poll::Ready(9), fut.as_mut().poll(&mut cx));
 154 }
 155
 156 fn main() {
 157     macro_rules! test {
 158         ($($fn_name:expr,)*) => { $(
 159             test_future_yields_once_then_returns($fn_name);
 160         )* }
 161     }
 162
 163     macro_rules! test_with_borrow {
 164         ($($fn_name:expr,)*) => { $(
 165             test_future_yields_once_then_returns(|x| {
 166                 async move {
 167                     $fn_name(&x).await
 168                 }
 169             });
 170         )* }
 171     }
 172
 173     test! {
 174         async_block,
 175         async_nonmove_block,
 176         async_fn,
 177         generic_async_fn,
 178         async_fn_with_internal_borrow,
 179         Foo::async_assoc_item,
 180         |x| {
 181             async move {
 182                 unsafe { unsafe_async_fn(x).await }
 183             }
 184         },
 185         |x| {
 186             async move {
 187                 unsafe { Foo::async_unsafe_assoc_item(x).await }
 188             }
 189         },
 190     }
 191     test_with_borrow! {
 192         async_block_with_borrow_named_lifetime,
 193         async_fn_with_borrow,
 194         async_fn_with_borrow_named_lifetime,
 195         async_fn_with_impl_future_named_lifetime,
 196         |x| {
 197             async move {
 198                 async_fn_multiple_args_named_lifetime(x, x).await
 199             }
 200         },
 201     }
 202 }