async_sink/ext/
err_into.rs

1use super::Sink;
2use core::marker::PhantomData;
3use core::pin::Pin;
4use core::task::{Context, Poll};
5use tokio_stream::Stream;
6
7/// Sink for the [`sink_err_into`](super::SinkExt::sink_err_into) method.
8#[derive(Debug)]
9#[must_use = "sinks do nothing unless polled"]
10pub struct SinkErrInto<Si, Item, E> {
11    sink: Si,
12    _phantom: PhantomData<fn(Item) -> E>,
13}
14
15impl<Si, Item, E> SinkErrInto<Si, Item, E> {
16    pub(super) fn new(sink: Si) -> Self {
17        Self {
18            sink,
19            _phantom: PhantomData,
20        }
21    }
22
23    /// Acquires a reference to the underlying sink.
24    pub fn get_ref(&self) -> &Si {
25        &self.sink
26    }
27
28    /// Acquires a mutable reference to the underlying sink.
29    ///
30    /// Note that care must be taken to avoid tampering with the state of the
31    /// sink which may otherwise confuse this combinator.
32    pub fn get_mut(&mut self) -> &mut Si {
33        &mut self.sink
34    }
35
36    /// Acquires a pinned mutable reference to the underlying sink.
37    ///
38    /// Note that care must be taken to avoid tampering with the state of the
39    /// sink which may otherwise confuse this combinator.
40    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut Si> {
41        unsafe { self.map_unchecked_mut(|s| &mut s.sink) }
42    }
43
44    /// Consumes this combinator, returning the underlying sink.
45    ///
46    /// Note that this may discard intermediate state of this combinator, so
47    /// care should be taken to avoid losing resources when this is called.
48    pub fn into_inner(self) -> Si {
49        self.sink
50    }
51}
52
53impl<Si, Item, E> Sink<Item> for SinkErrInto<Si, Item, E>
54where
55    Si: Sink<Item>,
56    Si::Error: Into<E>,
57    E: core::error::Error,
58{
59    type Error = E;
60
61    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
62        let sink = unsafe { self.map_unchecked_mut(|s| &mut s.sink) };
63        match sink.poll_ready(cx) {
64            Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
65            Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())),
66            Poll::Pending => Poll::Pending,
67        }
68    }
69
70    fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error> {
71        let sink = unsafe { self.map_unchecked_mut(|s| &mut s.sink) };
72        match sink.start_send(item) {
73            Ok(()) => Ok(()),
74            Err(e) => Err(e.into()),
75        }
76    }
77
78    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
79        let sink = unsafe { self.map_unchecked_mut(|s| &mut s.sink) };
80        match sink.poll_flush(cx) {
81            Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
82            Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())),
83            Poll::Pending => Poll::Pending,
84        }
85    }
86
87    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
88        let sink = unsafe { self.map_unchecked_mut(|s| &mut s.sink) };
89        match sink.poll_close(cx) {
90            Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
91            Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())),
92            Poll::Pending => Poll::Pending,
93        }
94    }
95}
96
97// Forwarding impl of Stream from the underlying sink
98impl<Si, Item, E> Stream for SinkErrInto<Si, Item, E>
99where
100    Si: Sink<Item> + Stream,
101{
102    type Item = Si::Item;
103
104    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
105        unsafe { self.map_unchecked_mut(|s| &mut s.sink) }.poll_next(cx)
106    }
107
108    fn size_hint(&self) -> (usize, Option<usize>) {
109        self.sink.size_hint()
110    }
111}