Implement FuturesUnordered::iter_mut

Author: srijs

src/stream/futures_unordered.rs

@@ -1,3 +1,5 @@
+//! An unbounded set of futures.
+
 use std::cell::UnsafeCell;
 use std::fmt::{self, Debug};
 use std::iter::FromIterator;
@@ -207,6 +209,15 @@ impl<T> FuturesUnordered<T> {
         self.inner.enqueue(ptr);
     }
 
+    /// Returns an iterator that allows modifying each future in the set.
+    pub fn iter_mut(&mut self) -> IterMut<T> {
+        IterMut {
+            node: self.head_all,
+            len: self.len,
+            _marker: PhantomData
+        }
+    }
+
     fn release_node(&mut self, node: Arc<Node<T>>) {
         // The future is done, try to reset the queued flag. This will prevent
         // `notify` from doing any work in the future
@@ -440,6 +451,37 @@ impl<F: Future> FromIterator<F> for FuturesUnordered<F> {
     }
 }
 
+#[derive(Debug)]
+/// Mutable iterator over all futures in the unordered set.
+pub struct IterMut<'a, F: 'a> {
+    node: *const Node<F>,
+    len: usize,
+    _marker: PhantomData<&'a F>
+}
+
+impl<'a, F> Iterator for IterMut<'a, F> {
+    type Item = &'a mut F;
+
+    fn next(&mut self) -> Option<&'a mut F> {
+        if self.node.is_null() {
+            return None;
+        }
+        unsafe {
+            let future = (*(*self.node).future.get()).as_mut().unwrap();
+            let next = *(*self.node).next_all.get();
+            self.node = next;
+            self.len -= 1;
+            return Some(future);
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.len, Some(self.len))
+    }
+}
+
+impl<'a, F> ExactSizeIterator for IterMut<'a, F> {}
+
 impl<T> Inner<T> {
     /// The enqueue function from the 1024cores intrusive MPSC queue algorithm.
     fn enqueue(&self, node: *const Node<T>) {

src/stream/mod.rs

@@ -103,7 +103,7 @@ if_std! {
     mod wait;
     mod channel;
     mod split;
-    mod futures_unordered;
+    pub mod futures_unordered;
     mod futures_ordered;
     pub use self::buffered::Buffered;
     pub use self::buffer_unordered::BufferUnordered;

tests/futures_unordered.rs

@@ -83,3 +83,45 @@ fn finished_future_ok() {
     assert!(spawn.poll_stream_notify(&support::notify_noop(), 0).unwrap().is_not_ready());
     assert!(spawn.poll_stream_notify(&support::notify_noop(), 0).unwrap().is_not_ready());
 }
+
+#[test]
+fn iter_mut_cancel() {
+    let (a_tx, a_rx) = oneshot::channel::<u32>();
+    let (b_tx, b_rx) = oneshot::channel::<u32>();
+    let (c_tx, c_rx) = oneshot::channel::<u32>();
+
+    let mut stream = futures_unordered(vec![a_rx, b_rx, c_rx]);
+
+    for rx in stream.iter_mut() {
+        rx.close();
+    }
+
+    assert!(a_tx.is_canceled());
+    assert!(b_tx.is_canceled());
+    assert!(c_tx.is_canceled());
+
+    let mut spawn = futures::executor::spawn(stream);
+    assert_eq!(Some(Err(futures::sync::oneshot::Canceled)), spawn.wait_stream());
+    assert_eq!(Some(Err(futures::sync::oneshot::Canceled)), spawn.wait_stream());
+    assert_eq!(Some(Err(futures::sync::oneshot::Canceled)), spawn.wait_stream());
+    assert_eq!(None, spawn.wait_stream());
+}
+
+#[test]
+fn iter_mut_len() {
+    let mut stream = futures_unordered(vec![
+        futures::future::empty::<(),()>(),
+        futures::future::empty::<(),()>(),
+        futures::future::empty::<(),()>()
+    ]);
+
+    let mut iter_mut = stream.iter_mut();
+    assert_eq!(iter_mut.len(), 3);
+    assert!(iter_mut.next().is_some());
+    assert_eq!(iter_mut.len(), 2);
+    assert!(iter_mut.next().is_some());
+    assert_eq!(iter_mut.len(), 1);
+    assert!(iter_mut.next().is_some());
+    assert_eq!(iter_mut.len(), 0);
+    assert!(iter_mut.next().is_none());
+}