rskit_stream/operators/
transform.rs

1use std::collections::HashSet;
2use std::hash::Hash;
3
4use futures::{Stream, StreamExt as _};
5use tokio::sync::mpsc;
6
7use crate::source;
8
9/// Emit only the first occurrence of each item.
10pub fn distinct<S, T>(stream: S) -> impl Stream<Item = T> + Send + 'static
11where
12    S: Stream<Item = T> + Send + 'static,
13    T: Clone + Eq + Hash + Send + 'static,
14{
15    async_stream::stream! {
16        tokio::pin!(stream);
17        let mut seen = HashSet::new();
18        while let Some(item) = stream.next().await {
19            if seen.insert(item.clone()) {
20                yield item;
21            }
22        }
23    }
24}
25
26/// Split a stream into two streams based on `predicate`.
27pub fn partition<S, T, F>(
28    stream: S,
29    mut predicate: F,
30) -> (
31    impl Stream<Item = T> + Send + 'static,
32    impl Stream<Item = T> + Send + 'static,
33)
34where
35    S: Stream<Item = T> + Send + 'static,
36    T: Send + 'static,
37    F: FnMut(&T) -> bool + Send + 'static,
38{
39    let (left_tx, left_rx) = mpsc::channel(64);
40    let (right_tx, right_rx) = mpsc::channel(64);
41
42    tokio::spawn(async move {
43        let mut left_open = true;
44        let mut right_open = true;
45
46        tokio::pin!(stream);
47        while let Some(item) = stream.next().await {
48            if !left_open && !right_open {
49                break;
50            }
51
52            if predicate(&item) {
53                if left_open && left_tx.send(item).await.is_err() {
54                    left_open = false;
55                }
56            } else if right_open && right_tx.send(item).await.is_err() {
57                right_open = false;
58            }
59        }
60    });
61
62    (
63        source::from_channel(left_rx),
64        source::from_channel(right_rx),
65    )
66}
rskit_stream/operators/transform.rs

rskit_stream/operators/
transform.rs
