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

   1 // run-pass
   2
   3 // edition:2018
   4 // aux-build:arc_wake.rs
   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(self: Arc<Self>) {
  23         Self::wake_by_ref(&self)
  24     }
  25     fn wake_by_ref(arc_self: &Arc<Self>) {
  26         arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
  27     }
  28 }
  29
  30 struct WakeOnceThenComplete(bool);
  31
  32 fn wake_and_yield_once() -> WakeOnceThenComplete { WakeOnceThenComplete(false) }
  33
  34 impl Future for WakeOnceThenComplete {
  35     type Output = ();
  36     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
  37         if self.0 {
  38             Poll::Ready(())
  39         } else {
  40             cx.waker().wake_by_ref();
  41             self.0 = true;
  42             Poll::Pending
  43         }
  44     }
  45 }
  46
  47 fn async_block(x: u8) -> impl Future<Output = u8> {
  48     async move {
  49         wake_and_yield_once().await;
  50         x
  51     }
  52 }
  53
  54 fn async_block_with_borrow_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
  55     async move {
  56         wake_and_yield_once().await;
  57         *x
  58     }
  59 }
  60
  61 fn async_nonmove_block(x: u8) -> impl Future<Output = u8> {
  62     async move {
  63         let future = async {
  64             wake_and_yield_once().await;
  65             x
  66         };
  67         future.await
  68     }
  69 }
  70
  71 // see async-closure.rs for async_closure + async_closure_in_unsafe_block
  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 async fn async_fn_multiple_args(x: &u8, _y: &u8) -> u8 {
 101     wake_and_yield_once().await;
 102     *x
 103 }
 104
 105 async fn async_fn_multiple_args_named_lifetime<'a>(x: &'a u8, _y: &'a u8) -> u8 {
 106     wake_and_yield_once().await;
 107     *x
 108 }
 109
 110 fn async_fn_with_internal_borrow(y: u8) -> impl Future<Output = u8> {
 111     async move {
 112         async_fn_with_borrow_named_lifetime(&y).await
 113     }
 114 }
 115
 116 async unsafe fn unsafe_async_fn(x: u8) -> u8 {
 117     wake_and_yield_once().await;
 118     x
 119 }
 120
 121 unsafe fn unsafe_fn(x: u8) -> u8 {
 122     x
 123 }
 124
 125 fn async_block_in_unsafe_block(x: u8) -> impl Future<Output = u8> {
 126     unsafe {
 127         async move {
 128             unsafe_fn(unsafe_async_fn(x).await)
 129         }
 130     }
 131 }
 132
 133 struct Foo;
 134
 135 trait Bar {
 136     fn foo() {}
 137 }
 138
 139 impl Foo {
 140     async fn async_assoc_item(x: u8) -> u8 {
 141         unsafe {
 142             unsafe_async_fn(x).await
 143         }
 144     }
 145
 146     async unsafe fn async_unsafe_assoc_item(x: u8) -> u8 {
 147         unsafe_async_fn(x).await
 148     }
 149 }
 150
 151 fn test_future_yields_once_then_returns<F, Fut>(f: F)
 152 where
 153     F: FnOnce(u8) -> Fut,
 154     Fut: Future<Output = u8>,
 155 {
 156     let mut fut = Box::pin(f(9));
 157     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
 158     let waker = ArcWake::into_waker(counter.clone());
 159     let mut cx = Context::from_waker(&waker);
 160     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
 161     assert_eq!(Poll::Pending, fut.as_mut().poll(&mut cx));
 162     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
 163     assert_eq!(Poll::Ready(9), fut.as_mut().poll(&mut cx));
 164 }
 165
 166 fn main() {
 167     macro_rules! test {
 168         ($($fn_name:expr,)*) => { $(
 169             test_future_yields_once_then_returns($fn_name);
 170         )* }
 171     }
 172
 173     macro_rules! test_with_borrow {
 174         ($($fn_name:expr,)*) => { $(
 175             test_future_yields_once_then_returns(|x| {
 176                 async move {
 177                     $fn_name(&x).await
 178                 }
 179             });
 180         )* }
 181     }
 182
 183     test! {
 184         async_block,
 185         async_nonmove_block,
 186         async_fn,
 187         generic_async_fn,
 188         async_fn_with_internal_borrow,
 189         async_block_in_unsafe_block,
 190         Foo::async_assoc_item,
 191         |x| {
 192             async move {
 193                 unsafe { unsafe_async_fn(x).await }
 194             }
 195         },
 196         |x| {
 197             async move {
 198                 unsafe { Foo::async_unsafe_assoc_item(x).await }
 199             }
 200         },
 201     }
 202     test_with_borrow! {
 203         async_block_with_borrow_named_lifetime,
 204         async_fn_with_borrow,
 205         async_fn_with_borrow_named_lifetime,
 206         async_fn_with_impl_future_named_lifetime,
 207         |x| {
 208             async move {
 209                 async_fn_multiple_args_named_lifetime(x, x).await
 210             }
 211         },
 212     }
 213 }