src/test/ui/auxiliary/issue-72470-lib.rs

   1 // compile-flags: -C opt-level=3
   2 // edition:2018
   3
   4 use std::future::Future;
   5 use std::marker::PhantomData;
   6 use std::pin::Pin;
   7 use std::sync::atomic::AtomicUsize;
   8 use std::sync::Arc;
   9 use std::task::Poll::{Pending, Ready};
  10 use std::task::Waker;
  11 use std::task::{Context, Poll};
  12 use std::{
  13     ptr,
  14     task::{RawWaker, RawWakerVTable},
  15 };
  16
  17 /// Future for the [`poll_fn`] function.
  18 pub struct PollFn<F> {
  19     f: F,
  20 }
  21
  22 impl<F> Unpin for PollFn<F> {}
  23
  24 /// Creates a new future wrapping around a function returning [`Poll`].
  25 pub fn poll_fn<T, F>(f: F) -> PollFn<F>
  26 where
  27     F: FnMut(&mut Context<'_>) -> Poll<T>,
  28 {
  29     PollFn { f }
  30 }
  31
  32 impl<T, F> Future for PollFn<F>
  33 where
  34     F: FnMut(&mut Context<'_>) -> Poll<T>,
  35 {
  36     type Output = T;
  37
  38     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
  39         (&mut self.f)(cx)
  40     }
  41 }
  42 pub fn run<F: Future>(future: F) -> F::Output {
  43     BasicScheduler.block_on(future)
  44 }
  45
  46 pub(crate) struct BasicScheduler;
  47
  48 impl BasicScheduler {
  49     pub(crate) fn block_on<F>(&mut self, mut future: F) -> F::Output
  50     where
  51         F: Future,
  52     {
  53         let waker = unsafe { Waker::from_raw(raw_waker()) };
  54         let mut cx = std::task::Context::from_waker(&waker);
  55
  56         let mut future = unsafe { Pin::new_unchecked(&mut future) };
  57
  58         loop {
  59             if let Ready(v) = future.as_mut().poll(&mut cx) {
  60                 return v;
  61             }
  62         }
  63     }
  64 }
  65
  66 // ===== impl Spawner =====
  67
  68 fn raw_waker() -> RawWaker {
  69     RawWaker::new(ptr::null(), waker_vtable())
  70 }
  71
  72 fn waker_vtable() -> &'static RawWakerVTable {
  73     &RawWakerVTable::new(
  74         clone_arc_raw,
  75         wake_arc_raw,
  76         wake_by_ref_arc_raw,
  77         drop_arc_raw,
  78     )
  79 }
  80
  81 unsafe fn clone_arc_raw(_: *const ()) -> RawWaker {
  82     raw_waker()
  83 }
  84
  85 unsafe fn wake_arc_raw(_: *const ()) {}
  86
  87 unsafe fn wake_by_ref_arc_raw(_: *const ()) {}
  88
  89 unsafe fn drop_arc_raw(_: *const ()) {}
  90
  91 struct AtomicWaker {}
  92
  93 impl AtomicWaker {
  94     /// Create an `AtomicWaker`
  95     fn new() -> AtomicWaker {
  96         AtomicWaker {}
  97     }
  98
  99     fn register_by_ref(&self, _waker: &Waker) {}
 100 }
 101
 102 #[allow(dead_code)]
 103 struct Tx<T> {
 104     inner: Arc<Chan<T>>,
 105 }
 106
 107 struct Rx<T> {
 108     inner: Arc<Chan<T>>,
 109 }
 110
 111 #[allow(dead_code)]
 112 struct Chan<T> {
 113     tx: PhantomData<T>,
 114     semaphore: Sema,
 115     rx_waker: AtomicWaker,
 116     rx_closed: bool,
 117 }
 118
 119 fn channel<T>() -> (Tx<T>, Rx<T>) {
 120     let chan = Arc::new(Chan {
 121         tx: PhantomData,
 122         semaphore: Sema(AtomicUsize::new(0)),
 123         rx_waker: AtomicWaker::new(),
 124         rx_closed: false,
 125     });
 126
 127     (
 128         Tx {
 129             inner: chan.clone(),
 130         },
 131         Rx { inner: chan },
 132     )
 133 }
 134
 135 // ===== impl Rx =====
 136
 137 impl<T> Rx<T> {
 138     /// Receive the next value
 139     fn recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> {
 140         self.inner.rx_waker.register_by_ref(cx.waker());
 141
 142         if self.inner.rx_closed && self.inner.semaphore.is_idle() {
 143             Ready(None)
 144         } else {
 145             Pending
 146         }
 147     }
 148 }
 149
 150 struct Sema(AtomicUsize);
 151
 152 impl Sema {
 153     fn is_idle(&self) -> bool {
 154         false
 155     }
 156 }
 157
 158 pub struct UnboundedReceiver<T> {
 159     chan: Rx<T>,
 160 }
 161
 162 pub fn unbounded_channel<T>() -> UnboundedReceiver<T> {
 163     let (tx, rx) = channel();
 164
 165     drop(tx);
 166     let rx = UnboundedReceiver { chan: rx };
 167
 168     rx
 169 }
 170
 171 impl<T> UnboundedReceiver<T> {
 172     pub async fn recv(&mut self) -> Option<T> {
 173         poll_fn(|cx| self.chan.recv(cx)).await
 174     }
 175 }