use crate::endpoint::CloseError;
use core::{
future::Future,
pin::Pin,
task::{Context, Poll},
};
use pin_project_lite::pin_project;
pin_project!(
pub struct Select<Rx, Tx, Wakeup, Sleep>
where
Rx: Future,
Tx: Future,
Wakeup: Future,
Sleep: Future,
{
#[pin]
rx: Rx,
#[pin]
tx: Tx,
#[pin]
wakeup: Wakeup,
#[pin]
sleep: Sleep,
}
);
impl<Rx, Tx, Wakeup, Sleep> Select<Rx, Tx, Wakeup, Sleep>
where
Rx: Future,
Tx: Future,
Wakeup: Future,
Sleep: Future,
{
#[inline(always)]
pub fn new(rx: Rx, tx: Tx, wakeup: Wakeup, sleep: Sleep) -> Self {
Self {
rx,
tx,
wakeup,
sleep,
}
}
}
#[derive(Debug)]
pub struct Outcome<Rx, Tx> {
pub rx_result: Option<Rx>,
pub tx_result: Option<Tx>,
pub timeout_expired: bool,
pub application_wakeup: bool,
}
pub type Result<Rx, Tx> = core::result::Result<Outcome<Rx, Tx>, CloseError>;
impl<Rx, Tx, Wakeup, Sleep> Future for Select<Rx, Tx, Wakeup, Sleep>
where
Rx: Future,
Tx: Future,
Wakeup: Future<Output = core::result::Result<usize, CloseError>>,
Sleep: Future,
{
type Output = Result<Rx::Output, Tx::Output>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
let mut should_wake = false;
let mut application_wakeup = false;
if let Poll::Ready(wakeup) = this.wakeup.poll(cx) {
should_wake = true;
application_wakeup = true;
if let Err(err) = wakeup {
return Poll::Ready(Err(err));
}
}
let mut rx_result = None;
if let Poll::Ready(v) = this.rx.poll(cx) {
should_wake = true;
rx_result = Some(v);
}
let mut tx_result = None;
if let Poll::Ready(v) = this.tx.poll(cx) {
should_wake = true;
tx_result = Some(v);
}
let mut timeout_expired = false;
if this.sleep.poll(cx).is_ready() {
timeout_expired = true;
should_wake = true;
}
if !should_wake {
return Poll::Pending;
}
Poll::Ready(Ok(Outcome {
rx_result,
tx_result,
timeout_expired,
application_wakeup,
}))
}
}