Merge pull request #12 from Palladinium/flat-map

Add flat_map

Author: sfackler

src/lib.rs

@@ -150,6 +150,22 @@ pub trait StreamingIterator {
         }
     }
 
+    /// Creates an iterator which flattens iterators obtained by applying a closure to elements.
+    /// Note that the returned iterators must be streaming iterators.
+    #[inline]
+    fn flat_map<J, F>(self, f: F) -> FlatMap<Self, J, F>
+    where
+        Self: Sized,
+        J: StreamingIterator,
+        F: FnMut(&Self::Item) -> J,
+    {
+        FlatMap {
+            it: self,
+            f,
+            sub_iter: None,
+        }
+    }
+
     /// Returns the first element of the iterator that satisfies the predicate.
     #[inline]
     fn find<F>(&mut self, mut f: F) -> Option<&Self::Item>
@@ -856,6 +872,40 @@ where
     }
 }
 
+/// A streaming iterator that maps elements to iterators with a closure and then yields the
+/// concatenation of the obtained iterators
+#[derive(Debug)]
+pub struct FlatMap<I, J, F> {
+    it: I,
+    f: F,
+    sub_iter: Option<J>,
+}
+
+impl<I, J, F> StreamingIterator for FlatMap<I, J, F>
+where
+    I: StreamingIterator,
+    F: FnMut(&I::Item) -> J,
+    J: StreamingIterator,
+{
+    type Item = J::Item;
+
+    #[inline]
+    fn advance(&mut self) {
+        while self.sub_iter.as_mut().and_then(J::next).is_none() {
+            if let Some(item) = self.it.next() {
+                self.sub_iter = Some((self.f)(item));
+            } else {
+                break;
+            }
+        }
+    }
+
+    #[inline]
+    fn get(&self) -> Option<&Self::Item> {
+        self.sub_iter.as_ref().and_then(J::get)
+    }
+}
+
 #[derive(Copy, Clone, Debug)]
 enum FuseState {
     Start,
@@ -1563,6 +1613,14 @@ mod test {
         test(it, &[0, 1]);
     }
 
+    #[test]
+    fn flat_map() {
+        let items = [[0, 1, 2], [3, 4, 5]];
+        let it = convert(items.iter()).flat_map(|i| convert(i.iter().cloned()));
+
+        test(it, &[0, 1, 2, 3, 4, 5]);
+    }
+
     #[test]
     fn nth() {
         let items = [0, 1];