pollable_map/stream/
set.rs

1use core::pin::Pin;
2use futures::{Stream, StreamExt};
3
4use super::StreamMap;
5use core::task::{Context, Poll};
6use futures::stream::FusedStream;
7
8pub struct StreamSet<S> {
9    id: i64,
10    map: StreamMap<i64, S>,
11}
12
13impl<S> Default for StreamSet<S>
14where
15    S: Stream,
16{
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl<S> StreamSet<S>
23where
24    S: Stream,
25{
26    /// Creates an empty [`StreamSet`]
27    pub fn new() -> Self {
28        Self {
29            id: 0,
30            map: StreamMap::default(),
31        }
32    }
33
34    /// Insert a stream into the set of streams.
35    pub fn insert(&mut self, stream: S) -> bool {
36        self.id = self.id.wrapping_add(1);
37        self.map.insert(self.id, stream)
38    }
39
40    /// An iterator visiting all streams in arbitrary order.
41    pub fn iter(&self) -> impl Iterator<Item = &S> {
42        self.map.iter().map(|(_, st)| st)
43    }
44
45    /// An iterator visiting all streams pinned valued in arbitrary order
46    pub fn iter_pin(&mut self) -> impl Iterator<Item = Pin<&mut S>> {
47        self.map.iter_pin().map(|(_, st)| st)
48    }
49
50    /// Clears the set.
51    pub fn clear(&mut self) {
52        self.map.clear();
53    }
54
55    /// Returns the number of streams in the set.
56    pub fn len(&self) -> usize {
57        self.map.len()
58    }
59
60    /// Return `true` map contains no elements.
61    pub fn is_empty(&self) -> bool {
62        self.map.is_empty()
63    }
64}
65
66impl<S> StreamSet<S>
67where
68    S: Stream + Unpin,
69{
70    /// An iterator visiting all streams mutably in arbitrary order.
71    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut S> {
72        self.map.iter_mut().map(|(_, st)| st)
73    }
74}
75
76impl<S> FromIterator<S> for StreamSet<S>
77where
78    S: Stream,
79{
80    fn from_iter<I: IntoIterator<Item = S>>(iter: I) -> Self {
81        let mut maps = Self::new();
82        for st in iter {
83            maps.insert(st);
84        }
85        maps
86    }
87}
88
89impl<S> Stream for StreamSet<S>
90where
91    S: Stream,
92{
93    type Item = S::Item;
94
95    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
96        self.map
97            .poll_next_unpin(cx)
98            .map(|output| output.map(|(_, item)| item))
99    }
100
101    fn size_hint(&self) -> (usize, Option<usize>) {
102        self.map.size_hint()
103    }
104}
105
106impl<S> FusedStream for StreamSet<S>
107where
108    S: Stream,
109{
110    fn is_terminated(&self) -> bool {
111        self.map.is_terminated()
112    }
113}
114
115#[cfg(test)]
116mod test {
117    use crate::stream::set::StreamSet;
118    use futures::StreamExt;
119
120    #[test]
121    fn valid_stream_set() {
122        let mut list = StreamSet::new();
123        assert!(list.insert(futures::stream::once(async { 0 }).boxed()));
124        assert!(list.insert(futures::stream::once(async { 1 }).boxed()));
125
126        futures::executor::block_on(async move {
127            let val = list.next().await;
128            assert_eq!(val, Some(0));
129            let val = list.next().await;
130            assert_eq!(val, Some(1));
131        });
132    }
133
134    #[test]
135    fn supports_unboxed_async_stream() {
136        let mut set = StreamSet::new();
137        let stream = futures::stream::unfold(0, |value| async move {
138            (value < 3).then_some((value, value + 1))
139        });
140        assert!(set.insert(stream));
141
142        futures::executor::block_on(async move {
143            assert_eq!(set.next().await, Some(0));
144            assert_eq!(set.next().await, Some(1));
145            assert_eq!(set.next().await, Some(2));
146            assert_eq!(set.next().await, None);
147        });
148    }
149}
pollable_map/stream/set.rs

pollable_map/stream/
set.rs
