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

   1 // run-pass
   2 // aux-build:arc_wake.rs
   3 // edition:2018
   4
   5 extern crate arc_wake;
   6
   7 use std::pin::Pin;
   8 use std::future::Future;
   9 use std::sync::{
  10     Arc,
  11     atomic::{self, AtomicUsize},
  12 };
  13 use std::task::{Context, Poll};
  14 use arc_wake::ArcWake;
  15
  16 struct Counter {
  17     wakes: AtomicUsize,
  18 }
  19
  20 impl ArcWake for Counter {
  21     fn wake(self: Arc<Self>) {
  22         Self::wake_by_ref(&self)
  23     }
  24     fn wake_by_ref(arc_self: &Arc<Self>) {
  25         arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
  26     }
  27 }
  28
  29 struct WakeOnceThenComplete(bool, u8);
  30
  31 impl Future for WakeOnceThenComplete {
  32     type Output = u8;
  33     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<u8> {
  34         if self.0 {
  35             Poll::Ready(self.1)
  36         } else {
  37             cx.waker().wake_by_ref();
  38             self.0 = true;
  39             Poll::Pending
  40         }
  41     }
  42 }
  43
  44 fn wait(fut: impl Future<Output = u8>) -> u8 {
  45     let mut fut = Box::pin(fut);
  46     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
  47     let waker = ArcWake::into_waker(counter.clone());
  48     let mut cx = Context::from_waker(&waker);
  49     loop {
  50         match fut.as_mut().poll(&mut cx) {
  51             Poll::Ready(out) => return out,
  52             Poll::Pending => (),
  53         }
  54     }
  55 }
  56
  57 fn base() -> WakeOnceThenComplete { WakeOnceThenComplete(false, 1) }
  58
  59 async fn await1_level1() -> u8 {
  60     base().await
  61 }
  62
  63 async fn await2_level1() -> u8 {
  64     base().await + base().await
  65 }
  66
  67 async fn await3_level1() -> u8 {
  68     base().await + base().await + base().await
  69 }
  70
  71 async fn await3_level2() -> u8 {
  72     await3_level1().await + await3_level1().await + await3_level1().await
  73 }
  74
  75 async fn await3_level3() -> u8 {
  76     await3_level2().await + await3_level2().await + await3_level2().await
  77 }
  78
  79 async fn await3_level4() -> u8 {
  80     await3_level3().await + await3_level3().await + await3_level3().await
  81 }
  82
  83 async fn await3_level5() -> u8 {
  84     await3_level4().await + await3_level4().await + await3_level4().await
  85 }
  86
  87 fn main() {
  88     assert_eq!(2, std::mem::size_of_val(&base()));
  89     assert_eq!(3, std::mem::size_of_val(&await1_level1()));
  90     assert_eq!(4, std::mem::size_of_val(&await2_level1()));
  91     assert_eq!(5, std::mem::size_of_val(&await3_level1()));
  92     assert_eq!(8, std::mem::size_of_val(&await3_level2()));
  93     assert_eq!(11, std::mem::size_of_val(&await3_level3()));
  94     assert_eq!(14, std::mem::size_of_val(&await3_level4()));
  95     assert_eq!(17, std::mem::size_of_val(&await3_level5()));
  96
  97     assert_eq!(1,   wait(base()));
  98     assert_eq!(1,   wait(await1_level1()));
  99     assert_eq!(2,   wait(await2_level1()));
 100     assert_eq!(3,   wait(await3_level1()));
 101     assert_eq!(9,   wait(await3_level2()));
 102     assert_eq!(27,  wait(await3_level3()));
 103     assert_eq!(81,  wait(await3_level4()));
 104     assert_eq!(243, wait(await3_level5()));
 105 }