add StreamingIteratorMut::flatten and Flatten

Author: LorenzSchueler

src/lib.rs

@@ -40,7 +40,7 @@
 #![warn(missing_docs)]
 #![cfg_attr(not(feature = "std"), no_std)]
 
-use core::cmp;
+use core::{cmp, marker::PhantomData};
 
 mod sources;
 pub use crate::sources::{convert, Convert};
@@ -592,6 +592,19 @@ pub trait StreamingIteratorMut: StreamingIterator {
     {
         MapDerefMut { it: self, f }
     }
+
+    /// Creates an iterator which flattens nested streaming iterators.
+    #[inline]
+    fn flatten(self) -> Flatten<Self, Self::Item>
+    where
+        Self: Sized,
+        Self::Item: StreamingIterator + Sized,
+    {
+        Flatten {
+            iter: self,
+            _inner: PhantomData,
+        }
+    }
 }
 
 impl<'a, I: ?Sized> StreamingIteratorMut for &'a mut I
@@ -1259,6 +1272,61 @@ where
     }
 }
 
+/// A streaming iterator that flattens nested streaming iterators.
+#[derive(Debug)]
+pub struct Flatten<I, J> {
+    iter: I,
+    _inner: PhantomData<J>,
+}
+
+impl<I, J> StreamingIterator for Flatten<I, J>
+where
+    I: StreamingIteratorMut<Item = J>,
+    J: StreamingIterator,
+{
+    type Item = J::Item;
+
+    #[inline]
+    fn advance(&mut self) {
+        loop {
+            if let Some(ref mut iter) = self.iter.get_mut() {
+                iter.advance();
+                if !iter.is_done() {
+                    break;
+                }
+            }
+            self.iter.advance();
+            if self.iter.is_done() {
+                break;
+            }
+        }
+    }
+
+    #[inline]
+    fn is_done(&self) -> bool {
+        match self.iter.get() {
+            Some(iter) => iter.is_done(),
+            None => true,
+        }
+    }
+
+    #[inline]
+    fn get(&self) -> Option<&Self::Item> {
+        self.iter.get().and_then(|iter| iter.get())
+    }
+}
+
+impl<I, J> StreamingIteratorMut for Flatten<I, J>
+where
+    I: StreamingIteratorMut<Item = J>,
+    J: StreamingIteratorMut,
+{
+    #[inline]
+    fn get_mut(&mut self) -> Option<&mut Self::Item> {
+        self.iter.get_mut().and_then(J::get_mut)
+    }
+}
+
 /// A regular, non-streaming iterator which both filters and maps elements of a streaming iterator with a closure.
 #[derive(Debug)]
 pub struct FilterMapDeref<I, F> {