use tokio::sync::mpsc;
#[derive(Debug)]
pub struct PeekableReceiver<T> {
rx: mpsc::Receiver<T>,
peeked: Option<T>,
}
impl<T> From<mpsc::Receiver<T>> for PeekableReceiver<T> {
fn from(rx: mpsc::Receiver<T>) -> Self {
Self::new(rx)
}
}
impl<T> PeekableReceiver<T> {
pub fn new(rx: mpsc::Receiver<T>) -> Self {
Self { rx, peeked: None }
}
#[allow(dead_code)]
pub async fn recv(&mut self) -> Option<T> {
match self.peeked.take() {
Some(msg) => Some(msg),
None => self.rx.recv().await,
}
}
pub fn try_recv(&mut self) -> Result<T, mpsc::error::TryRecvError> {
match self.peeked.take() {
Some(msg) => Ok(msg),
None => self.rx.try_recv(),
}
}
pub async fn peek(&mut self) -> Option<&T> {
if self.peeked.is_none() {
self.peeked = self.rx.recv().await;
}
self.peeked.as_ref()
}
pub fn try_peek(&mut self) -> Result<&T, mpsc::error::TryRecvError> {
if self.peeked.is_none() {
self.peeked = Some(self.rx.try_recv()?);
}
Ok(self.peeked.as_ref().unwrap())
}
pub async fn recv_if(&mut self, cond: impl FnOnce(&T) -> bool) -> Result<T, RecvIfError> {
match self.peek().await {
Some(msg) if cond(msg) => Ok(self.try_recv().unwrap()),
Some(_) => Err(RecvIfError::NoMatch),
None => Err(RecvIfError::Disconnected),
}
}
pub fn try_recv_if(&mut self, cond: impl FnOnce(&T) -> bool) -> Result<T, TryRecvIfError> {
if cond(self.try_peek()?) {
Ok(self.try_recv().unwrap())
} else {
Err(TryRecvIfError::NoMatch)
}
}
}
#[derive(Debug, Clone)]
pub enum RecvIfError {
Disconnected,
NoMatch,
}
#[derive(Debug, Clone)]
pub enum TryRecvIfError {
Empty,
Disconnected,
NoMatch,
}
impl From<mpsc::error::TryRecvError> for TryRecvIfError {
fn from(err: mpsc::error::TryRecvError) -> Self {
match err {
mpsc::error::TryRecvError::Empty => Self::Empty,
mpsc::error::TryRecvError::Disconnected => Self::Disconnected,
}
}
}