src/test/run-pass/async-await.rs

   1 // Copyright 2018 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 // edition:2018
  12
  13 #![feature(arbitrary_self_types, async_await, await_macro, futures_api, pin)]
  14
  15 use std::pin::Pin;
  16 use std::future::Future;
  17 use std::sync::{
  18     Arc,
  19     atomic::{self, AtomicUsize},
  20 };
  21 use std::task::{
  22     LocalWaker, Poll, Wake,
  23     local_waker_from_nonlocal,
  24 };
  25
  26 struct Counter {
  27     wakes: AtomicUsize,
  28 }
  29
  30 impl Wake for Counter {
  31     fn wake(this: &Arc<Self>) {
  32         this.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>, lw: &LocalWaker) -> Poll<()> {
  43         if self.0 {
  44             Poll::Ready(())
  45         } else {
  46             lw.wake();
  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         await!(wake_and_yield_once());
  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         await!(wake_and_yield_once());
  63         *x
  64     }
  65 }
  66
  67 fn async_nonmove_block(x: u8) -> impl Future<Output = u8> {
  68     async move {
  69         let future = async {
  70             await!(wake_and_yield_once());
  71             x
  72         };
  73         await!(future)
  74     }
  75 }
  76
  77 fn async_closure(x: u8) -> impl Future<Output = u8> {
  78     (async move |x: u8| -> u8 {
  79         await!(wake_and_yield_once());
  80         x
  81     })(x)
  82 }
  83
  84 async fn async_fn(x: u8) -> u8 {
  85     await!(wake_and_yield_once());
  86     x
  87 }
  88
  89 async fn async_fn_with_borrow(x: &u8) -> u8 {
  90     await!(wake_and_yield_once());
  91     *x
  92 }
  93
  94 async fn async_fn_with_borrow_named_lifetime<'a>(x: &'a u8) -> u8 {
  95     await!(wake_and_yield_once());
  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         await!(wake_and_yield_once());
 102         *x
 103     }
 104 }
 105
 106 async fn async_fn_with_named_lifetime_multiple_args<'a>(x: &'a u8, _y: &'a u8) -> u8 {
 107     await!(wake_and_yield_once());
 108     *x
 109 }
 110
 111 fn async_fn_with_internal_borrow(y: u8) -> impl Future<Output = u8> {
 112     async move {
 113         await!(async_fn_with_borrow(&y))
 114     }
 115 }
 116
 117 unsafe async fn unsafe_async_fn(x: u8) -> u8 {
 118     await!(wake_and_yield_once());
 119     x
 120 }
 121
 122 struct Foo;
 123
 124 trait Bar {
 125     fn foo() {}
 126 }
 127
 128 impl Foo {
 129     async fn async_method(x: u8) -> u8 {
 130         unsafe {
 131             await!(unsafe_async_fn(x))
 132         }
 133     }
 134 }
 135
 136 fn test_future_yields_once_then_returns<F, Fut>(f: F)
 137 where
 138     F: FnOnce(u8) -> Fut,
 139     Fut: Future<Output = u8>,
 140 {
 141     let mut fut = Box::pinned(f(9));
 142     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
 143     let waker = local_waker_from_nonlocal(counter.clone());
 144     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
 145     assert_eq!(Poll::Pending, fut.as_mut().poll(&waker));
 146     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
 147     assert_eq!(Poll::Ready(9), fut.as_mut().poll(&waker));
 148 }
 149
 150 fn main() {
 151     macro_rules! test {
 152         ($($fn_name:expr,)*) => { $(
 153             test_future_yields_once_then_returns($fn_name);
 154         )* }
 155     }
 156
 157     macro_rules! test_with_borrow {
 158         ($($fn_name:expr,)*) => { $(
 159             test_future_yields_once_then_returns(|x| {
 160                 async move {
 161                     await!($fn_name(&x))
 162                 }
 163             });
 164         )* }
 165     }
 166
 167     test! {
 168         async_block,
 169         async_nonmove_block,
 170         async_closure,
 171         async_fn,
 172         async_fn_with_internal_borrow,
 173         |x| {
 174             async move {
 175                 unsafe { await!(unsafe_async_fn(x)) }
 176             }
 177         },
 178     }
 179
 180     test_with_borrow! {
 181         async_block_with_borrow_named_lifetime,
 182         async_fn_with_borrow,
 183         async_fn_with_borrow_named_lifetime,
 184         async_fn_with_impl_future_named_lifetime,
 185         |x| {
 186             async move {
 187                 await!(async_fn_with_named_lifetime_multiple_args(x, x))
 188             }
 189         },
 190     }
 191 }