ruchei_extend/
keyed.rs

1extern crate std;
2
3use core::{
4    pin::Pin,
5    task::{Context, Poll},
6};
7
8use extend_pinned::ExtendPinned;
9#[cfg(feature = "sink")]
10use futures_sink::Sink;
11use futures_util::{Stream, stream::FusedStream};
12use pin_project::pin_project;
13#[cfg(feature = "route-sink")]
14use route_sink::{FlushRoute, ReadyRoute, ReadySome};
15
16use crate::{Extending, ExtendingExt};
17
18pub enum ExtendItem<K, T> {
19    Pushed(K),
20    Item(T),
21}
22
23impl<K, T, E> ExtendItem<K, Result<T, E>> {
24    pub fn transpose(self) -> Result<ExtendItem<K, T>, E> {
25        match self {
26            Self::Pushed(x) => Ok(ExtendItem::Pushed(x)),
27            Self::Item(Ok(x)) => Ok(ExtendItem::Item(x)),
28            Self::Item(Err(x)) => Err(x),
29        }
30    }
31}
32
33#[pin_project]
34#[derive(Debug)]
35pub struct ExtendKeyed<S, K> {
36    #[pin]
37    inner: S,
38    pushed: std::collections::VecDeque<K>,
39}
40
41impl<S: Default, K> Default for ExtendKeyed<S, K> {
42    fn default() -> Self {
43        S::default().into()
44    }
45}
46
47impl<S, K> From<S> for ExtendKeyed<S, K> {
48    fn from(inner: S) -> Self {
49        Self {
50            inner,
51            pushed: Default::default(),
52        }
53    }
54}
55
56impl<S: ExtendPinned<T>, K, T> ExtendPinned<(K, T)> for ExtendKeyed<S, K> {
57    fn extend_pinned<I: IntoIterator<Item = (K, T)>>(self: Pin<&mut Self>, iter: I) {
58        let this = self.project();
59        this.inner.extend_pinned(iter.into_iter().map(|(k, t)| {
60            this.pushed.push_back(k);
61            t
62        }));
63    }
64}
65
66impl<S: Stream, K> Stream for ExtendKeyed<S, K> {
67    type Item = ExtendItem<K, S::Item>;
68
69    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
70        let this = self.project();
71        if let Some(k) = this.pushed.pop_front() {
72            return Poll::Ready(Some(ExtendItem::Pushed(k)));
73        }
74        this.inner.poll_next(cx).map(|o| o.map(ExtendItem::Item))
75    }
76}
77
78impl<S: FusedStream, K> FusedStream for ExtendKeyed<S, K> {
79    fn is_terminated(&self) -> bool {
80        self.pushed.is_empty() && self.inner.is_terminated()
81    }
82}
83
84#[cfg(feature = "sink")]
85impl<Item, S: Sink<Item>, K> Sink<Item> for ExtendKeyed<S, K> {
86    type Error = S::Error;
87
88    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
89        self.project().inner.poll_ready(cx)
90    }
91
92    fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error> {
93        self.project().inner.start_send(item)
94    }
95
96    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
97        self.project().inner.poll_flush(cx)
98    }
99
100    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
101        self.project().inner.poll_close(cx)
102    }
103}
104
105#[cfg(feature = "route-sink")]
106impl<Route, Msg, S: FlushRoute<Route, Msg>, K> FlushRoute<Route, Msg> for ExtendKeyed<S, K> {
107    fn poll_flush_route(
108        self: Pin<&mut Self>,
109        route: &Route,
110        cx: &mut Context<'_>,
111    ) -> Poll<Result<(), Self::Error>> {
112        self.project().inner.poll_flush_route(route, cx)
113    }
114
115    fn poll_close_route(
116        self: Pin<&mut Self>,
117        route: &Route,
118        cx: &mut Context<'_>,
119    ) -> Poll<Result<(), Self::Error>> {
120        self.project().inner.poll_close_route(route, cx)
121    }
122}
123
124#[cfg(feature = "route-sink")]
125impl<Route, Msg, S: ReadyRoute<Route, Msg>, K> ReadyRoute<Route, Msg> for ExtendKeyed<S, K> {
126    fn poll_ready_route(
127        self: Pin<&mut Self>,
128        route: &Route,
129        cx: &mut Context<'_>,
130    ) -> Poll<Result<(), Self::Error>> {
131        self.project().inner.poll_ready_route(route, cx)
132    }
133}
134
135#[cfg(feature = "route-sink")]
136impl<Route, Msg, S: ReadySome<Route, Msg>, K> ReadySome<Route, Msg> for ExtendKeyed<S, K> {
137    fn poll_ready_some(
138        self: Pin<&mut Self>,
139        cx: &mut Context<'_>,
140    ) -> Poll<Result<Route, Self::Error>> {
141        self.project().inner.poll_ready_some(cx)
142    }
143}
144
145pub trait ExtendingKeyedExt: Sized + Stream<Item = (Self::K, Self::T)> {
146    type K;
147    type T;
148
149    #[must_use]
150    fn extending_keyed<S: ExtendPinned<Self::T>>(
151        self,
152        inner: S,
153    ) -> Extending<ExtendKeyed<S, Self::K>, Self> {
154        self.extending(inner.into())
155    }
156
157    #[must_use]
158    fn extending_keyed_default<S: Default + ExtendPinned<Self::T>>(
159        self,
160    ) -> Extending<ExtendKeyed<S, Self::K>, Self> {
161        self.extending_default()
162    }
163}
ruchei_extend/keyed.rs

ruchei_extend/
keyed.rs
