specialize Intersperse::fold

Author: jswrenn

benches/fold_specialization.rs

@@ -0,0 +1,66 @@
+#![feature(test)]
+
+extern crate test;
+extern crate itertools;
+
+use itertools::Itertools;
+
+struct Unspecialized<I>(I);
+
+impl<I> Iterator for Unspecialized<I>
+where I: Iterator
+{
+    type Item = I::Item;
+
+    #[inline(always)]
+    fn next(&mut self) -> Option<I::Item> {
+        self.0.next()
+    }
+
+    #[inline(always)]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.0.size_hint()
+    }
+}
+
+mod specialization {
+    use super::*;
+
+    mod intersperse {
+        use super::*;
+
+        #[bench]
+        fn external(b: &mut test::Bencher)
+        {
+            let arr = [1; 1024];
+
+            b.iter(|| {
+                let mut sum = 0;
+                for &x in arr.into_iter().intersperse(&0) {
+                    sum += x;
+                }
+                sum
+            })
+        }
+
+        #[bench]
+        fn internal_specialized(b: &mut test::Bencher)
+        {
+            let arr = [1; 1024];
+
+            b.iter(|| {
+                arr.into_iter().intersperse(&0).fold(0, |acc, x| acc + x)
+            })
+        }
+
+        #[bench]
+        fn internal_unspecialized(b: &mut test::Bencher)
+        {
+            let arr = [1; 1024];
+
+            b.iter(|| {
+                Unspecialized(arr.into_iter().intersperse(&0)).fold(0, |acc, x| acc + x)
+            })
+        }
+    }
+}

src/intersperse.rs

@@ -57,4 +57,23 @@ impl<I> Iterator for Intersperse<I>
         let sh = self.iter.size_hint();
         size_hint::add_scalar(size_hint::add(sh, sh), has_peek)
     }
+
+    fn fold<B, F>(mut self, init: B, mut f: F) -> B where
+        Self: Sized, F: FnMut(B, Self::Item) -> B,
+    {
+        let mut accum = init;
+        
+        if let Some(x) = self.peek.take() {
+            accum = f(accum, x);
+        }
+
+        let element = &self.element;
+
+        self.iter.fold(accum,
+            |accum, x| {
+                let accum = f(accum, element.clone());
+                let accum = f(accum, x);
+                accum
+        })
+    }
 }

tests/fold_specialization.rs

@@ -0,0 +1,15 @@
+extern crate itertools;
+
+use itertools::Itertools;
+
+#[test]
+fn specialization_intersperse() {
+    let mut iter = (1..2).intersperse(0);
+    iter.clone().for_each(|x| assert_eq!(Some(x), iter.next()));
+
+    let mut iter = (1..3).intersperse(0);
+    iter.clone().for_each(|x| assert_eq!(Some(x), iter.next()));
+
+    let mut iter = (1..4).intersperse(0);
+    iter.clone().for_each(|x| assert_eq!(Some(x), iter.next()));
+}