juggle/
yield_helper.rs

1use core::future::Future;
2use core::pin::Pin;
3use core::task::{Context, Poll};
4
5/// Helper struct for dealing with task switching.
6#[derive(Clone, Debug, Hash, Eq, PartialEq)]
7pub struct Yield(bool);
8
9#[doc(hidden)]
10#[derive(Clone, Debug, Hash, Eq, PartialEq)]
11pub struct YieldWhile<F: FnMut() -> bool>(F);
12
13#[doc(hidden)]
14#[derive(Clone, Debug, Hash, Eq, PartialEq)]
15pub struct YieldTimes { pub remaining: usize }
16
17impl Yield {
18    /// When awaited yields this task once. Causes task switch.
19    ///
20    /// For more convenient method of switching tasks see [`yield_once!()`](macro.yield_once.html) macro.
21    ///
22    /// When resulting Future is polled, for the first time it notifies the waker and returns
23    /// `Poll::Pending`, second and all other polls return `Poll::Ready(())`.
24    pub fn once() -> Self { Self(false) }
25
26    /// When awaited it won't cause task switch.
27    ///
28    /// Future returned by this method when polled always return `Poll::Ready(())`.
29    pub fn none() -> Self { Self(true) }
30
31    /// When awaited yields this task specific number of times.
32    ///
33    /// Resulting Future notifies the waker and returns
34    /// `Poll::Pending` 'remaining' number of times, all other polls return `Poll::Ready(())`.
35    pub fn times(remaining: usize) -> YieldTimes { YieldTimes { remaining } }
36
37    /// When awaited yields this task until provided closure returns false.
38    ///
39    /// Note that when first call on closure returns false, this task will not be yielded.
40    /// This method is usefull when we want to do busy wait but also leave cpu time for
41    /// other tasks.
42    /// # Examples
43    /// ```
44    /// # fn main(){
45    /// # use juggle::Yield;
46    /// # use core::sync::atomic::{AtomicBool, Ordering};
47    /// # smol::block_on(async move{
48    /// let interrupt_flag: &AtomicBool = //...
49    /// # &AtomicBool::new(true);
50    ///
51    /// Yield::yield_while(||!interrupt_flag.load(Ordering::Acquire)).await;
52    /// # });
53    /// # }
54    /// ```
55    pub fn yield_while<F>(predicate: F) -> YieldWhile<F> where F: FnMut() -> bool {
56        YieldWhile(predicate)
57    }
58}
59
60impl Future for Yield {
61    type Output = ();
62    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
63        if self.0 { Poll::Ready(()) } else {
64            self.get_mut().0 = true;
65            cx.waker().wake_by_ref();
66            Poll::Pending
67        }
68    }
69}
70
71impl<F: FnMut() -> bool> Future for YieldWhile<F> {
72    type Output = ();
73    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
74        //SAFETY: F does not care about being pinned
75        let func = unsafe{ &mut self.get_unchecked_mut().0 };
76        if !func() { Poll::Ready(()) } else {
77            cx.waker().wake_by_ref();
78            Poll::Pending
79        }
80    }
81}
82
83impl Future for YieldTimes {
84    type Output = ();
85    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
86        if self.remaining == 0 { Poll::Ready(()) } else {
87            self.as_mut().remaining -= 1;
88            cx.waker().wake_by_ref();
89            Poll::Pending
90        }
91    }
92}
juggle/yield_helper.rs

juggle/
yield_helper.rs
