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

   1 // run-pass
   2
   3 #![allow(unused)]
   4
   5 // edition:2018
   6 // aux-build:arc_wake.rs
   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 // see async-closure.rs for async_closure + async_closure_in_unsafe_block
  74
  75 async fn async_fn(x: u8) -> u8 {
  76     wake_and_yield_once().await;
  77     x
  78 }
  79
  80 async fn generic_async_fn<T>(x: T) -> T {
  81     wake_and_yield_once().await;
  82     x
  83 }
  84
  85 async fn async_fn_with_borrow(x: &u8) -> u8 {
  86     wake_and_yield_once().await;
  87     *x
  88 }
  89
  90 async fn async_fn_with_borrow_named_lifetime<'a>(x: &'a u8) -> u8 {
  91     wake_and_yield_once().await;
  92     *x
  93 }
  94
  95 fn async_fn_with_impl_future_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
  96     async move {
  97         wake_and_yield_once().await;
  98         *x
  99     }
 100 }
 101
 102 async fn async_fn_multiple_args(x: &u8, _y: &u8) -> u8 {
 103     wake_and_yield_once().await;
 104     *x
 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 unsafe fn unsafe_fn(x: u8) -> u8 {
 124     x
 125 }
 126
 127 fn async_block_in_unsafe_block(x: u8) -> impl Future<Output = u8> {
 128     unsafe {
 129         async move {
 130             unsafe_fn(unsafe_async_fn(x).await)
 131         }
 132     }
 133 }
 134
 135 struct Foo;
 136
 137 trait Bar {
 138     fn foo() {}
 139 }
 140
 141 impl Foo {
 142     async fn async_assoc_item(x: u8) -> u8 {
 143         unsafe {
 144             unsafe_async_fn(x).await
 145         }
 146     }
 147
 148     async unsafe fn async_unsafe_assoc_item(x: u8) -> u8 {
 149         unsafe_async_fn(x).await
 150     }
 151 }
 152
 153 fn test_future_yields_once_then_returns<F, Fut>(f: F)
 154 where
 155     F: FnOnce(u8) -> Fut,
 156     Fut: Future<Output = u8>,
 157 {
 158     let mut fut = Box::pin(f(9));
 159     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
 160     let waker = ArcWake::into_waker(counter.clone());
 161     let mut cx = Context::from_waker(&waker);
 162     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
 163     assert_eq!(Poll::Pending, fut.as_mut().poll(&mut cx));
 164     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
 165     assert_eq!(Poll::Ready(9), fut.as_mut().poll(&mut cx));
 166 }
 167
 168 fn main() {
 169     macro_rules! test {
 170         ($($fn_name:expr,)*) => { $(
 171             test_future_yields_once_then_returns($fn_name);
 172         )* }
 173     }
 174
 175     macro_rules! test_with_borrow {
 176         ($($fn_name:expr,)*) => { $(
 177             test_future_yields_once_then_returns(|x| {
 178                 async move {
 179                     $fn_name(&x).await
 180                 }
 181             });
 182         )* }
 183     }
 184
 185     test! {
 186         async_block,
 187         async_nonmove_block,
 188         async_fn,
 189         generic_async_fn,
 190         async_fn_with_internal_borrow,
 191         async_block_in_unsafe_block,
 192         Foo::async_assoc_item,
 193         |x| {
 194             async move {
 195                 unsafe { unsafe_async_fn(x).await }
 196             }
 197         },
 198         |x| {
 199             async move {
 200                 unsafe { Foo::async_unsafe_assoc_item(x).await }
 201             }
 202         },
 203     }
 204     test_with_borrow! {
 205         async_block_with_borrow_named_lifetime,
 206         async_fn_with_borrow,
 207         async_fn_with_borrow_named_lifetime,
 208         async_fn_with_impl_future_named_lifetime,
 209         |x| {
 210             async move {
 211                 async_fn_multiple_args_named_lifetime(x, x).await
 212             }
 213         },
 214     }
 215 }