Implement IterMut in terms of IterPinMut

Author: MajorBreakfast

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

@@ -1,33 +1,21 @@
-use std::marker::{PhantomData, Unpin};
-use super::FuturesUnordered;
-use super::node::Node;
+use std::marker::Unpin;
+use std::mem::PinMut;
+
+use super::iter_pin_mut::IterPinMut;
 
 #[derive(Debug)]
 /// Mutable iterator over all futures in the unordered set.
-pub struct IterMut<'a, F: 'a + Unpin> {
-    pub(super) node: *const Node<F>,
-    pub(super) len: usize,
-    pub(super) _marker: PhantomData<&'a mut FuturesUnordered<F>>
-}
+pub struct IterMut<'a, F: 'a + Unpin> (pub(super) IterPinMut<'a, F>);
 
 impl<'a, F: Unpin> 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;
-            Some(future)
-        }
+        self.0.next().map(|f| unsafe { PinMut::get_mut(f) })
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        (self.len, Some(self.len))
+        self.0.size_hint()
     }
 }
 

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

@@ -1,4 +1,4 @@
-use std::marker::PhantomData;
+use std::marker::{PhantomData, Unpin};
 use std::mem::PinMut;
 
 use super::FuturesUnordered;
@@ -12,6 +12,8 @@ pub struct IterPinMut<'a, F: 'a> {
     pub(super) _marker: PhantomData<&'a mut FuturesUnordered<F>>
 }
 
+impl<'a, F: Unpin> Unpin for IterPinMut<'a, F> {}
+
 impl<'a, F> Iterator for IterPinMut<'a, F> {
     type Item = PinMut<'a, F>;
 

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

@@ -150,11 +150,11 @@ impl<T> FuturesUnordered<T> {
 
     /// Returns an iterator that allows modifying each future in the set.
     pub fn iter_mut(&mut self) -> IterMut<T> where T: Unpin {
-        IterMut {
+        IterMut(IterPinMut {
             node: self.head_all,
             len: self.len,
             _marker: PhantomData
-        }
+        })
     }
 
     /// Returns an iterator that allows modifying each future in the set.