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

   1 // run-pass
   2
   3 // revisions: default nomiropt
   4 //[nomiropt]compile-flags: -Z mir-opt-level=0
   5
   6 #![allow(unused)]
   7
   8 // edition:2018
   9 // aux-build:arc_wake.rs
  10
  11 extern crate arc_wake;
  12
  13 use std::pin::Pin;
  14 use std::future::Future;
  15 use std::sync::{
  16     Arc,
  17     atomic::{self, AtomicUsize},
  18 };
  19 use std::task::{Context, Poll};
  20 use arc_wake::ArcWake;
  21
  22 struct Counter {
  23     wakes: AtomicUsize,
  24 }
  25
  26 impl ArcWake for Counter {
  27     fn wake(self: Arc<Self>) {
  28         Self::wake_by_ref(&self)
  29     }
  30     fn wake_by_ref(arc_self: &Arc<Self>) {
  31         arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
  32     }
  33 }
  34
  35 struct WakeOnceThenComplete(bool);
  36
  37 fn wake_and_yield_once() -> WakeOnceThenComplete { WakeOnceThenComplete(false) }
  38
  39 impl Future for WakeOnceThenComplete {
  40     type Output = ();
  41     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
  42         if self.0 {
  43             Poll::Ready(())
  44         } else {
  45             cx.waker().wake_by_ref();
  46             self.0 = true;
  47             Poll::Pending
  48         }
  49     }
  50 }
  51
  52 fn async_block(x: u8) -> impl Future<Output = u8> {
  53     async move {
  54         wake_and_yield_once().await;
  55         x
  56     }
  57 }
  58
  59 fn async_block_with_borrow_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
  60     async move {
  61         wake_and_yield_once().await;
  62         *x
  63     }
  64 }
  65
  66 fn async_nonmove_block(x: u8) -> impl Future<Output = u8> {
  67     async move {
  68         let future = async {
  69             wake_and_yield_once().await;
  70             x
  71         };
  72         future.await
  73     }
  74 }
  75
  76 // see async-closure.rs for async_closure + async_closure_in_unsafe_block
  77
  78 async fn async_fn(x: u8) -> u8 {
  79     wake_and_yield_once().await;
  80     x
  81 }
  82
  83 async fn generic_async_fn<T>(x: T) -> T {
  84     wake_and_yield_once().await;
  85     x
  86 }
  87
  88 async fn async_fn_with_borrow(x: &u8) -> u8 {
  89     wake_and_yield_once().await;
  90     *x
  91 }
  92
  93 async fn async_fn_with_borrow_named_lifetime<'a>(x: &'a u8) -> u8 {
  94     wake_and_yield_once().await;
  95     *x
  96 }
  97
  98 fn async_fn_with_impl_future_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
  99     async move {
 100         wake_and_yield_once().await;
 101         *x
 102     }
 103 }
 104
 105 async fn async_fn_multiple_args(x: &u8, _y: &u8) -> u8 {
 106     wake_and_yield_once().await;
 107     *x
 108 }
 109
 110 async fn async_fn_multiple_args_named_lifetime<'a>(x: &'a u8, _y: &'a u8) -> u8 {
 111     wake_and_yield_once().await;
 112     *x
 113 }
 114
 115 fn async_fn_with_internal_borrow(y: u8) -> impl Future<Output = u8> {
 116     async move {
 117         async_fn_with_borrow_named_lifetime(&y).await
 118     }
 119 }
 120
 121 async unsafe fn unsafe_async_fn(x: u8) -> u8 {
 122     wake_and_yield_once().await;
 123     x
 124 }
 125
 126 unsafe fn unsafe_fn(x: u8) -> u8 {
 127     x
 128 }
 129
 130 fn async_block_in_unsafe_block(x: u8) -> impl Future<Output = u8> {
 131     unsafe {
 132         async move {
 133             unsafe_fn(unsafe_async_fn(x).await)
 134         }
 135     }
 136 }
 137
 138 struct Foo;
 139
 140 trait Bar {
 141     fn foo() {}
 142 }
 143
 144 impl Foo {
 145     async fn async_assoc_item(x: u8) -> u8 {
 146         unsafe {
 147             unsafe_async_fn(x).await
 148         }
 149     }
 150
 151     async unsafe fn async_unsafe_assoc_item(x: u8) -> u8 {
 152         unsafe_async_fn(x).await
 153     }
 154 }
 155
 156 fn test_future_yields_once_then_returns<F, Fut>(f: F)
 157 where
 158     F: FnOnce(u8) -> Fut,
 159     Fut: Future<Output = u8>,
 160 {
 161     let mut fut = Box::pin(f(9));
 162     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
 163     let waker = ArcWake::into_waker(counter.clone());
 164     let mut cx = Context::from_waker(&waker);
 165     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
 166     assert_eq!(Poll::Pending, fut.as_mut().poll(&mut cx));
 167     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
 168     assert_eq!(Poll::Ready(9), fut.as_mut().poll(&mut cx));
 169 }
 170
 171 fn main() {
 172     macro_rules! test {
 173         ($($fn_name:expr,)*) => { $(
 174             test_future_yields_once_then_returns($fn_name);
 175         )* }
 176     }
 177
 178     macro_rules! test_with_borrow {
 179         ($($fn_name:expr,)*) => { $(
 180             test_future_yields_once_then_returns(|x| {
 181                 async move {
 182                     $fn_name(&x).await
 183                 }
 184             });
 185         )* }
 186     }
 187
 188     test! {
 189         async_block,
 190         async_nonmove_block,
 191         async_fn,
 192         generic_async_fn,
 193         async_fn_with_internal_borrow,
 194         async_block_in_unsafe_block,
 195         Foo::async_assoc_item,
 196         |x| {
 197             async move {
 198                 unsafe { unsafe_async_fn(x).await }
 199             }
 200         },
 201         |x| {
 202             async move {
 203                 unsafe { Foo::async_unsafe_assoc_item(x).await }
 204             }
 205         },
 206     }
 207     test_with_borrow! {
 208         async_block_with_borrow_named_lifetime,
 209         async_fn_with_borrow,
 210         async_fn_with_borrow_named_lifetime,
 211         async_fn_with_impl_future_named_lifetime,
 212         |x| {
 213             async move {
 214                 async_fn_multiple_args_named_lifetime(x, x).await
 215             }
 216         },
 217     }
 218 }