src/tools/miri/tests/pass/stacked-borrows/future-self-referential.rs

   1 use std::future::*;
   2 use std::marker::PhantomPinned;
   3 use std::pin::*;
   4 use std::ptr;
   5 use std::task::*;
   6
   7 struct Delay {
   8     delay: usize,
   9 }
  10
  11 impl Delay {
  12     fn new(delay: usize) -> Self {
  13         Delay { delay }
  14     }
  15 }
  16
  17 impl Future for Delay {
  18     type Output = ();
  19     fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
  20         if self.delay > 0 {
  21             self.delay -= 1;
  22             Poll::Pending
  23         } else {
  24             Poll::Ready(())
  25         }
  26     }
  27 }
  28
  29 fn mk_waker() -> Waker {
  30     use std::sync::Arc;
  31
  32     struct MyWaker;
  33     impl Wake for MyWaker {
  34         fn wake(self: Arc<Self>) {
  35             unimplemented!()
  36         }
  37     }
  38
  39     Waker::from(Arc::new(MyWaker))
  40 }
  41
  42 async fn do_stuff() {
  43     (&mut Delay::new(1)).await;
  44 }
  45
  46 // Same thing implemented by hand
  47 struct DoStuff {
  48     state: usize,
  49     delay: Delay,
  50     delay_ref: *mut Delay,
  51     _marker: PhantomPinned,
  52 }
  53
  54 impl DoStuff {
  55     fn new() -> Self {
  56         DoStuff {
  57             state: 0,
  58             delay: Delay::new(1),
  59             delay_ref: ptr::null_mut(),
  60             _marker: PhantomPinned,
  61         }
  62     }
  63 }
  64
  65 impl Future for DoStuff {
  66     type Output = ();
  67     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
  68         unsafe {
  69             let this = self.get_unchecked_mut();
  70             match this.state {
  71                 0 => {
  72                     // Set up self-ref.
  73                     this.delay_ref = &mut this.delay;
  74                     // Move to next state.
  75                     this.state = 1;
  76                     Poll::Pending
  77                 }
  78                 1 => {
  79                     let delay = &mut *this.delay_ref;
  80                     Pin::new_unchecked(delay).poll(cx)
  81                 }
  82                 _ => unreachable!(),
  83             }
  84         }
  85     }
  86 }
  87
  88 fn run_fut<T>(fut: impl Future<Output = T>) -> T {
  89     let waker = mk_waker();
  90     let mut context = Context::from_waker(&waker);
  91
  92     let mut pinned = pin!(fut);
  93     loop {
  94         match pinned.as_mut().poll(&mut context) {
  95             Poll::Pending => continue,
  96             Poll::Ready(v) => return v,
  97         }
  98     }
  99 }
 100
 101 fn self_referential_box() {
 102     let waker = mk_waker();
 103     let cx = &mut Context::from_waker(&waker);
 104
 105     async fn my_fut() -> i32 {
 106         let val = 10;
 107         let val_ref = &val;
 108
 109         let _ = Delay::new(1).await;
 110
 111         *val_ref
 112     }
 113
 114     fn box_poll<F: Future>(
 115         mut f: Pin<Box<F>>,
 116         cx: &mut Context<'_>,
 117     ) -> (Pin<Box<F>>, Poll<F::Output>) {
 118         let p = f.as_mut().poll(cx);
 119         (f, p)
 120     }
 121
 122     let my_fut = Box::pin(my_fut());
 123     let (my_fut, p1) = box_poll(my_fut, cx);
 124     assert!(p1.is_pending());
 125     let (my_fut, p2) = box_poll(my_fut, cx);
 126     assert!(p2.is_ready());
 127     drop(my_fut);
 128 }
 129
 130 fn main() {
 131     run_fut(do_stuff());
 132     run_fut(DoStuff::new());
 133     self_referential_box();
 134 }