async_let/
wait.rs

1use core::{
2    future::Future,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use crate::{
8    list::{At, Empty, FutList},
9    ReadyOrNot,
10};
11
12pub(crate) use private::DriveWaitFor;
13
14pin_project_lite::pin_project! {
15    /// Future type for the [`Group::wait_for`] method.
16    ///
17    /// [`wait_for`]: super::Group::wait_for
18    #[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
19    pub struct WaitFor<'group, F, List> {
20        #[pin]
21        pub(crate) driving_fut: F,
22        pub(crate) async_let_group: &'group mut List,
23    }
24}
25
26impl<F: Future, List: FutList> Future for WaitFor<'_, F, List> {
27    type Output = F::Output;
28
29    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
30        let this = self.project();
31        let poll = this.driving_fut.poll(cx);
32        if poll.is_pending() {
33            this.async_let_group.poll_once(cx);
34        }
35        poll
36    }
37}
38
39mod private {
40    /// Helper trait to poll each async let future when a waited on future is polled.
41    pub trait DriveWaitFor {
42        fn poll_once(&mut self, cx: &mut super::Context<'_>);
43    }
44}
45
46impl DriveWaitFor for Empty {
47    #[inline]
48    fn poll_once(&mut self, _cx: &mut Context<'_>) {}
49}
50
51impl<F: Future + Unpin, T: DriveWaitFor> DriveWaitFor for At<F, T>
52{
53    fn poll_once(&mut self, cx: &mut Context<'_>) {
54        let At { node, tail, .. } = self;
55        if let ReadyOrNot::Not(fut) = node {
56            if let Poll::Ready(val) = Pin::new(fut).poll(cx) {
57                *node = ReadyOrNot::Ready(val);
58            }
59        }
60        tail.poll_once(cx);
61    }
62}