1
//! Extension trait for `Sink`.
2

            
3
use std::{
4
    marker::PhantomData,
5
    pin::Pin,
6
    task::{Context, Poll},
7
};
8

            
9
use futures::{ready, sink::Sink};
10
use pin_project::pin_project;
11

            
12
/// Extension trait for `Sink`
13
pub trait SinkExt<Item>: Sink<Item> {
14
    /// As `Sink::with`, but takes a function that returns an `Item` rather
15
    /// than `Future<Output=Item>`.
16
4
    fn with_fn<F, T, E>(self, func: F) -> WithFn<Self, F, T, E>
17
4
    // or error?
18
4
    where
19
4
        Self: Sized,
20
4
        F: FnMut(T) -> Result<Item, E>,
21
4
        E: From<Self::Error>,
22
    {
23
4
        WithFn {
24
4
            sink: self,
25
4
            func,
26
4
            _phantom: PhantomData,
27
4
        }
28
4
    }
29
}
30

            
31
impl<Item, S> SinkExt<Item> for S where S: Sink<Item> {}
32

            
33
/// Sink returned by [`SinkExt::with_fn`].
34
#[pin_project]
35
pub struct WithFn<S, F, T, E> {
36
    /// The underlying sink
37
    #[pin]
38
    sink: S,
39
    /// The user-provided function.
40
    func: F,
41
    /// Phantom data to ensure type consistency.
42
    _phantom: PhantomData<fn() -> Result<T, E>>,
43
}
44

            
45
impl<S, Item, F, T, E> Sink<T> for WithFn<S, F, T, E>
46
where
47
    S: Sink<Item>,
48
    F: FnMut(T) -> Result<Item, E>,
49
    E: From<S::Error>,
50
{
51
    type Error = E;
52

            
53
4
    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
54
4
        ready!(self.project().sink.poll_ready(cx))?;
55
4
        Poll::Ready(Ok(()))
56
4
    }
57

            
58
4
    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
59
4
        ready!(self.project().sink.poll_flush(cx))?;
60
4
        Poll::Ready(Ok(()))
61
4
    }
62

            
63
    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
64
        ready!(self.project().sink.poll_close(cx))?;
65
        Poll::Ready(Ok(()))
66
    }
67

            
68
4
    fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
69
4
        let this = self.project();
70
4
        let item = (this.func)(item)?;
71
4
        this.sink.start_send(item).map_err(E::from)
72
4
    }
73
}