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 async fn do_stuff() {
  30     (&mut Delay::new(1)).await;
  31 }
  32
  33 // Same thing implemented by hand
  34 struct DoStuff {
  35     state: usize,
  36     delay: Delay,
  37     delay_ref: *mut Delay,
  38     _marker: PhantomPinned,
  39 }
  40
  41 impl DoStuff {
  42     fn new() -> Self {
  43         DoStuff {
  44             state: 0,
  45             delay: Delay::new(1),
  46             delay_ref: ptr::null_mut(),
  47             _marker: PhantomPinned,
  48         }
  49     }
  50 }
  51
  52 impl Future for DoStuff {
  53     type Output = ();
  54     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
  55         unsafe {
  56             let this = self.get_unchecked_mut();
  57             match this.state {
  58                 0 => {
  59                     // Set up self-ref.
  60                     this.delay_ref = &mut this.delay;
  61                     // Move to next state.
  62                     this.state = 1;
  63                     Poll::Pending
  64                 }
  65                 1 => {
  66                     let delay = &mut *this.delay_ref;
  67                     Pin::new_unchecked(delay).poll(cx)
  68                 }
  69                 _ => unreachable!(),
  70             }
  71         }
  72     }
  73 }
  74
  75 fn run_fut<T>(fut: impl Future<Output = T>) -> T {
  76     use std::sync::Arc;
  77
  78     struct MyWaker;
  79     impl Wake for MyWaker {
  80         fn wake(self: Arc<Self>) {
  81             unimplemented!()
  82         }
  83     }
  84
  85     let waker = Waker::from(Arc::new(MyWaker));
  86     let mut context = Context::from_waker(&waker);
  87
  88     let mut pinned = pin!(fut);
  89     loop {
  90         match pinned.as_mut().poll(&mut context) {
  91             Poll::Pending => continue,
  92             Poll::Ready(v) => return v,
  93         }
  94     }
  95 }
  96
  97 fn main() {
  98     run_fut(do_stuff());
  99     run_fut(DoStuff::new());
 100 }