Add extension to filter Result<T, E> by predicate(&T)

1 files changed, 62 insertions, 0 deletions

diff --git a/src/iter.rs b/src/iter.rs
index 79675d1..359625b 100644
--- a/src/iter.rs
+++ b/src/iter.rs
@@ -39,6 +39,47 @@ impl<I, T, F: Filter<T>> FilterWith<T, F> for I
     }
 }
 
+
+pub struct FilterOksIter<T, E, I, F>(I, F)
+    where F: Filter<T>,
+          I: Iterator<Item = Result<T, E>>;
+
+impl<T, E, I, F> Iterator for FilterOksIter<T, E, I, F>
+    where F: Filter<T>,
+          I: Iterator<Item = Result<T, E>>
+{
+    type Item = Result<T, E>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        while let Some(next) = self.0.next() {
+            match next {
+                Err(e) => return Some(Err(e)),
+                Ok(t)  => if self.1.filter(&t) {
+                    return Some(Ok(t));
+                }
+            }
+        }
+
+        None
+    }
+}
+
+pub trait FilterOks<T, E, I, F> : Iterator<Item = Result<T, E>>
+    where I: Iterator<Item = Result<T, E>>,
+          F: Filter<T>
+{
+    fn filter_oks(self, F) -> FilterOksIter<T, E, I, F>;
+}
+
+impl<T, E, I, F> FilterOks<T, E, I, F> for I
+    where I: Iterator<Item = Result<T, E>>,
+          F: Filter<T>
+{
+    fn filter_oks(self, f: F) -> FilterOksIter<T, E, I, F> {
+        FilterOksIter(self, f)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -61,4 +102,25 @@ mod test {
 
         assert_eq!(v, vec![6, 7, 8, 9]);
     }
+
+    #[test]
+    fn test_filter_oks() {
+        struct Foo;
+        impl Filter<u64> for Foo {
+            fn filter(&self, u: &u64) -> bool {
+                *u > 5
+            }
+        }
+
+        let foo = Foo;
+
+        let v : Vec<Result<u64, u64>> =
+            vec![Ok(1), Err(2), Ok(3), Err(4), Ok(5), Err(6), Ok(7), Err(8), Ok(9), Err(0)]
+            .into_iter()
+            .filter_oks(foo)
+            .collect();
+
+        assert_eq!(v, vec![Err(2), Err(4), Err(6), Ok(7), Err(8), Ok(9), Err(0)]);
+    }
 }
+