1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

use core::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};

/// A future that yields to the scheduler one or more times before completing.
#[derive(Debug)]
#[must_use = "futures do nothing unless `.await`ed or polled"]
pub struct Yield {
    yields: usize,
}

impl Future for Yield {
    type Output = ();
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        let yields = &mut self.as_mut().yields;
        if *yields == 0 {
            return Poll::Ready(());
        }
        *yields -= 1;
        cx.waker().wake_by_ref();
        Poll::Pending
    }
}

impl Yield {
    /// Returns a new future that yields `yields` times before completing.
    #[inline]
    pub const fn new(yields: usize) -> Self {
        Self { yields }
    }
}

/// Yield to the scheduler a single time before proceeding.
#[inline]
pub fn yield_now() -> Yield {
    Yield::new(1)
}