1 // run-pass
2
3 // revisions: default nomiropt
4 //[nomiropt]compile-flags: -Z mir-opt-level=0
5
6 // edition:2018
7 // aux-build:arc_wake.rs
8
9 #![feature(async_closure)]
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 {
wake(self: Arc<Self>)27 fn wake(self: Arc<Self>) {
28 Self::wake_by_ref(&self)
29 }
wake_by_ref(arc_self: &Arc<Self>)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
wake_and_yield_once() -> WakeOnceThenComplete37 fn wake_and_yield_once() -> WakeOnceThenComplete { WakeOnceThenComplete(false) }
38
39 impl Future for WakeOnceThenComplete {
40 type Output = ();
poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()>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
async_closure(x: u8) -> impl Future<Output = u8>52 fn async_closure(x: u8) -> impl Future<Output = u8> {
53 (async move |x: u8| -> u8 {
54 wake_and_yield_once().await;
55 x
56 })(x)
57 }
58
async_closure_in_unsafe_block(x: u8) -> impl Future<Output = u8>59 fn async_closure_in_unsafe_block(x: u8) -> impl Future<Output = u8> {
60 (unsafe {
61 async move |x: u8| unsafe_fn(unsafe_async_fn(x).await)
62 })(x)
63 }
64
unsafe_async_fn(x: u8) -> u865 async unsafe fn unsafe_async_fn(x: u8) -> u8 {
66 wake_and_yield_once().await;
67 x
68 }
69
unsafe_fn(x: u8) -> u870 unsafe fn unsafe_fn(x: u8) -> u8 {
71 x
72 }
73
test_future_yields_once_then_returns<F, Fut>(f: F) where F: FnOnce(u8) -> Fut, Fut: Future<Output = u8>,74 fn test_future_yields_once_then_returns<F, Fut>(f: F)
75 where
76 F: FnOnce(u8) -> Fut,
77 Fut: Future<Output = u8>,
78 {
79 let mut fut = Box::pin(f(9));
80 let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
81 let waker = ArcWake::into_waker(counter.clone());
82 let mut cx = Context::from_waker(&waker);
83 assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
84 assert_eq!(Poll::Pending, fut.as_mut().poll(&mut cx));
85 assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
86 assert_eq!(Poll::Ready(9), fut.as_mut().poll(&mut cx));
87 }
88
main()89 fn main() {
90 macro_rules! test {
91 ($($fn_name:expr,)*) => { $(
92 test_future_yields_once_then_returns($fn_name);
93 )* }
94 }
95
96 test! {
97 async_closure,
98 async_closure_in_unsafe_block,
99 }
100 }
101