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

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