Restrict FuturesUnorderdered::iter_mut to Unpin futures. Add iter_pin_mut

Author: MajorBreakfast

futures-util/src/stream/futures_unordered/iter_mut.rs

@@ -1,4 +1,4 @@
-use std::marker::PhantomData;
+use std::marker::{PhantomData, Unpin};
 use super::FuturesUnordered;
 use super::node::Node;
 
@@ -10,7 +10,7 @@ pub struct IterMut<'a, F: 'a> {
     pub(super) _marker: PhantomData<&'a mut FuturesUnordered<F>>
 }
 
-impl<'a, F> Iterator for IterMut<'a, F> {
+impl<'a, F: Unpin> Iterator for IterMut<'a, F> {
     type Item = &'a mut F;
 
     fn next(&mut self) -> Option<&'a mut F> {
@@ -31,4 +31,4 @@ impl<'a, F> Iterator for IterMut<'a, F> {
     }
 }
 
-impl<'a, F> ExactSizeIterator for IterMut<'a, F> {}
+impl<'a, F: Unpin> ExactSizeIterator for IterMut<'a, F> {}

futures-util/src/stream/futures_unordered/iter_pin_mut.rs

@@ -0,0 +1,36 @@
+use std::marker::PhantomData;
+use std::mem::PinMut;
+
+use super::FuturesUnordered;
+use super::node::Node;
+
+#[derive(Debug)]
+/// Mutable iterator over all futures in the unordered set.
+pub struct IterPinMut<'a, F: 'a> {
+    pub(super) node: *const Node<F>,
+    pub(super) len: usize,
+    pub(super) _marker: PhantomData<&'a mut FuturesUnordered<F>>
+}
+
+impl<'a, F> Iterator for IterPinMut<'a, F> {
+    type Item = PinMut<'a, F>;
+
+    fn next(&mut self) -> Option<PinMut<'a, 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;
+            Some(PinMut::new_unchecked(future))
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.len, Some(self.len))
+    }
+}
+
+impl<'a, F> ExactSizeIterator for IterPinMut<'a, F> {}

futures-util/src/stream/futures_unordered/mod.rs

@@ -17,10 +17,12 @@ use futures_core::task::{self, AtomicWaker};
 mod abort;
 mod ready_to_run_queue;
 mod iter_mut;
+mod iter_pin_mut;
 mod node;
 
 use self::ready_to_run_queue::{ReadyToRunQueue, Dequeue};
 use self::iter_mut::IterMut;
+use self::iter_pin_mut::IterPinMut;
 use self::node::Node;
 
 /// A set of `Future`s which may complete in any order.
@@ -155,6 +157,15 @@ impl<T> FuturesUnordered<T> {
         }
     }
 
+    /// Returns an iterator that allows modifying each future in the set.
+    pub fn iter_pin_mut<'a>(self: PinMut<'a, Self>) -> IterPinMut<'a, T> {
+        IterPinMut {
+            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