use core::pin::Pin;
use pin_project_lite::pin_project;
use crate::{Sink, SinkBuild, forward_sink};
pin_project! {
#[must_use = "sinks do nothing unless polled"]
pub struct Filter<Si, Func> {
#[pin]
sink: Si,
func: Func,
}
}
impl<Si, Func> Filter<Si, Func> {
pub fn new<Item>(func: Func, sink: Si) -> Self
where
Self: Sink<Item>,
{
Self { sink, func }
}
}
impl<Si, Func, Item> Sink<Item> for Filter<Si, Func>
where
Si: Sink<Item>,
Func: FnMut(&Item) -> bool,
{
type Error = Si::Error;
fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error> {
let this = self.project();
if (this.func)(&item) {
this.sink.start_send(item)
} else {
Ok(())
}
}
forward_sink!(poll_ready, poll_flush, poll_close);
}
pub struct FilterBuilder<Prev, Func> {
pub(crate) prev: Prev,
pub(crate) func: Func,
}
impl<Prev, Func> SinkBuild for FilterBuilder<Prev, Func>
where
Prev: SinkBuild,
Func: FnMut(&Prev::Item) -> bool,
{
type Item = Prev::Item;
type Output<Next: Sink<Prev::Item>> = Prev::Output<Filter<Next, Func>>;
fn send_to<Next>(self, next: Next) -> Self::Output<Next>
where
Next: Sink<Prev::Item>,
{
self.prev.send_to(Filter::new(self.func, next))
}
}