pollable_map/
futures.rs

1pub mod ordered;
2pub mod set;
3#[cfg(all(feature = "std", feature = "timeout"))]
4pub mod timeout_map;
5#[cfg(all(feature = "std", feature = "timeout"))]
6pub mod timeout_set;
7
8use crate::common::InnerMap;
9use core::future::Future;
10use core::pin::Pin;
11use core::task::{Context, Poll, Waker};
12use futures::stream::{FusedStream, FuturesUnordered};
13use futures::{Stream, StreamExt};
14
15pub struct FutureMap<K, S> {
16    list: FuturesUnordered<InnerMap<K, S>>,
17    empty: bool,
18    terminate_on_empty: bool,
19    waker: Option<Waker>,
20}
21
22impl<K, T> Default for FutureMap<K, T> {
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28impl<K, T> FutureMap<K, T> {
29    /// Creates an empty [`FutureMap`]
30    pub fn new() -> Self {
31        Self {
32            list: FuturesUnordered::new(),
33            empty: true,
34            terminate_on_empty: false,
35            waker: None,
36        }
37    }
38
39    /// Set flag to terminate stream after all futures are completed
40    pub fn set_terminate_on_empty(&mut self, terminate: bool) {
41        self.terminate_on_empty = terminate;
42    }
43}
44
45impl<K, T> FutureMap<K, T>
46where
47    K: Clone + PartialEq,
48    T: Future,
49{
50    /// Insert a future into the map with a unique key.
51    /// The function will return true if the map does not have the key present,
52    /// otherwise it will return false
53    pub fn insert(&mut self, key: K, fut: T) -> bool {
54        if self.contains_key(&key) {
55            return false;
56        }
57
58        let st = InnerMap::new(key, fut);
59        self.list.push(st);
60
61        if let Some(waker) = self.waker.take() {
62            waker.wake();
63        }
64
65        self.empty = false;
66        true
67    }
68
69    /// Mark future with assigned key to wake up on successful yield.
70    /// Will return false if future does not exist or if value is the same as
71    /// previously set.
72    pub fn set_wake_on_success(&mut self, key: &K, wake_on_success: bool) -> bool {
73        Pin::new(&mut self.list)
74            .iter_pin_mut()
75            .find(|st| st.as_ref().key_pin().eq(key))
76            .is_some_and(|st| st.set_wake_on_success_pin(wake_on_success))
77    }
78
79    /// An iterator visiting all key-value pairs in arbitrary order.
80    pub fn iter(&self) -> impl Iterator<Item = (&K, &T)> {
81        Pin::new(&self.list)
82            .iter_pin_ref()
83            .filter_map(|st| st.key_value_pin_ref())
84            .map(|(key, future)| (key, future.get_ref()))
85    }
86
87    /// An iterator visiting all key-value pairs with a pinned valued in arbitrary order
88    pub fn iter_pin(&mut self) -> impl Iterator<Item = (&K, Pin<&mut T>)> {
89        Pin::new(&mut self.list)
90            .iter_pin_mut()
91            .filter_map(|st| st.key_value_pin())
92    }
93
94    /// Returns an iterator visiting all keys in arbitrary order.
95    pub fn keys(&self) -> impl Iterator<Item = &K> {
96        Pin::new(&self.list)
97            .iter_pin_ref()
98            .filter_map(|st| st.key_value_pin_ref().map(|(key, _)| key))
99    }
100
101    /// An iterator visiting all values in arbitrary order.
102    pub fn values(&self) -> impl Iterator<Item = &T> {
103        Pin::new(&self.list)
104            .iter_pin_ref()
105            .filter_map(|st| st.inner_pin_ref())
106            .map(Pin::get_ref)
107    }
108
109    /// Returns `true` if the map contains a future for the specified key.
110    pub fn contains_key(&self, key: &K) -> bool {
111        Pin::new(&self.list)
112            .iter_pin_ref()
113            .filter(|st| st.as_ref().inner_pin_ref().is_some())
114            .any(|st| st.key_pin().eq(key))
115    }
116
117    /// Clears the map.
118    pub fn clear(&mut self) {
119        self.list.clear();
120    }
121
122    /// Returns a reference to the future corresponding to the key.
123    pub fn get(&self, key: &K) -> Option<&T> {
124        Pin::new(&self.list)
125            .iter_pin_ref()
126            .find(|st| st.as_ref().key_pin().eq(key))
127            .and_then(|st| st.inner_pin_ref())
128            .map(Pin::get_ref)
129    }
130
131    /// Returns a pinned future corresponding to the key.
132    pub fn get_pinned(&mut self, key: &K) -> Option<Pin<&mut T>> {
133        Pin::new(&mut self.list)
134            .iter_pin_mut()
135            .find(|st| st.as_ref().key_pin().eq(key))
136            .and_then(|st| st.inner_pin())
137    }
138
139    /// Returns the number of futures in the map.
140    pub fn len(&self) -> usize {
141        Pin::new(&self.list)
142            .iter_pin_ref()
143            .filter(|st| st.as_ref().inner_pin_ref().is_some())
144            .count()
145    }
146
147    /// Return `true` map contains no elements.
148    pub fn is_empty(&self) -> bool {
149        self.len() == 0
150    }
151}
152
153impl<K, T> FutureMap<K, T>
154where
155    K: Clone + PartialEq,
156    T: Future + Unpin,
157{
158    /// An iterator visiting all key-value pairs mutably in arbitrary order.
159    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&K, &mut T)> {
160        self.list.iter_mut().filter_map(|st| st.key_value_mut())
161    }
162
163    /// An iterator visiting all values mutably in arbitrary order.
164    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut T> {
165        self.list.iter_mut().filter_map(|st| st.inner_mut())
166    }
167
168    /// Returns a mutable future corresponding to the key.
169    pub fn get_mut(&mut self, key: &K) -> Option<&mut T> {
170        self.list
171            .iter_mut()
172            .find(|st| st.key().eq(key))
173            .and_then(|st| st.inner_mut())
174    }
175
176    /// Returns a mutable future or default value if it does not exist.
177    pub fn get_mut_or_default(&mut self, key: &K) -> &mut T
178    where
179        T: Default,
180    {
181        self.insert(key.clone(), T::default());
182        self.get_mut(key).expect("valid entry")
183    }
184
185    /// Removes a key from the map, returning the future.
186    pub fn remove(&mut self, key: &K) -> Option<T> {
187        self.list
188            .iter_mut()
189            .find(|st| st.key().eq(key))
190            .and_then(|st| st.take_inner())
191    }
192}
193
194impl<K, T> FromIterator<(K, T)> for FutureMap<K, T>
195where
196    K: Clone + PartialEq,
197    T: Future,
198{
199    fn from_iter<I: IntoIterator<Item = (K, T)>>(iter: I) -> Self {
200        let mut maps = Self::new();
201        for (key, val) in iter {
202            maps.insert(key, val);
203        }
204        maps
205    }
206}
207
208impl<K, T> Stream for FutureMap<K, T>
209where
210    K: Clone,
211    T: Future,
212{
213    type Item = (K, T::Output);
214
215    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
216        loop {
217            match self.list.poll_next_unpin(cx) {
218                Poll::Ready(Some((key, Some(item)))) => return Poll::Ready(Some((key, item))),
219                // We continue in case there is any progress on the set of streams
220                Poll::Ready(Some((_key, None))) => continue,
221                Poll::Ready(None) => {
222                    // While we could allow the stream to continue to be pending, it would make more sense to notify that the stream
223                    // is empty without needing to explicitly check while polling the actual "map" itself
224                    // So we would mark a field to notify that the state is finished and return `Poll::Ready(None)` so the stream
225                    // can be terminated while on the next poll, we could let it be return pending.
226                    // We do this so that we are not returning `Poll::Ready(None)` each time the map is polled
227                    // as that may be seen as UB and may cause an increase in cpu usage
228                    if self.empty {
229                        if self.terminate_on_empty {
230                            return Poll::Ready(None);
231                        }
232                        self.waker = Some(cx.waker().clone());
233                        return Poll::Pending;
234                    }
235
236                    self.empty = true;
237                    return Poll::Ready(None);
238                }
239                Poll::Pending => {
240                    // Returning `None` does not mean the stream is actually terminated
241                    self.waker = Some(cx.waker().clone());
242                    return Poll::Pending;
243                }
244            }
245        }
246    }
247
248    fn size_hint(&self) -> (usize, Option<usize>) {
249        self.list.size_hint()
250    }
251}
252
253impl<K, T> FusedStream for FutureMap<K, T>
254where
255    K: Clone,
256    T: Future,
257{
258    fn is_terminated(&self) -> bool {
259        self.terminate_on_empty && self.list.is_terminated()
260    }
261}
262
263#[cfg(test)]
264mod test {
265    use crate::futures::FutureMap;
266    use core::task::Poll;
267    use futures::future::pending;
268    use futures::StreamExt;
269
270    #[test]
271    fn existing_key() {
272        let mut map = FutureMap::new();
273        assert!(map.insert(1, pending::<()>()));
274        assert!(!map.insert(1, pending::<()>()));
275    }
276
277    #[test]
278    fn supports_unboxed_async_future() {
279        let mut map = FutureMap::new();
280        assert!(map.insert(1, async { 42 }));
281
282        futures::executor::block_on(async move {
283            assert_eq!(map.next().await, Some((1, 42)));
284        });
285    }
286
287    #[test]
288    fn poll_multiple_keyed_streams() {
289        let mut map = FutureMap::new();
290        map.insert(1, futures::future::ready(10));
291        map.insert(2, futures::future::ready(20));
292        map.insert(3, futures::future::ready(30));
293
294        futures::executor::block_on(async move {
295            assert_eq!(map.next().await, Some((1, 10)));
296            assert_eq!(map.next().await, Some((2, 20)));
297            assert_eq!(map.next().await, Some((3, 30)));
298            assert_eq!(map.next().await, None);
299            let pending =
300                futures::future::poll_fn(|cx| Poll::Ready(map.poll_next_unpin(cx).is_pending()))
301                    .await;
302            assert!(pending);
303        })
304    }
305}
pollable_map/futures.rs

pollable_map/
futures.rs
