Repository: bluss/rust-itertools
Pull request #269 (2018-03-18 01:49:48)

Add Itertools::tree_fold1
This is useful for non-associative operations like building a tree (folding over a binop constructor) or combining floating point data (where this order keeps the magnitudes more similar).


Patch:
diff --git a/.travis.yml b/.travis.yml
index ae9acd4..5bd2734 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -18,9 +18,6 @@ matrix:
 branches:
   only:
     - master
-    # bors branches
-    - staging
-    - trying
 script:
   - |
       cargo build --verbose --no-default-features &&
diff --git a/benches/tree_fold1.rs b/benches/tree_fold1.rs
new file mode 100644
index 0000000..b71589f
--- /dev/null
+++ b/benches/tree_fold1.rs
@@ -0,0 +1,126 @@
+#![feature(test)]
+
+extern crate test;
+extern crate itertools;
+
+use itertools::Itertools;
+use itertools::cloned;
+use test::Bencher;
+
+trait IterEx : Iterator {
+    // Another efficient implementation against which to compare,
+    // but needs `std` so is less desirable.
+    fn tree_fold1_vec<F>(self, mut f: F) -> Option<Self::Item>
+        where F: FnMut(Self::Item, Self::Item) -> Self::Item,
+              Self: Sized,
+    {
+        let hint = self.size_hint().0;
+        let cap = std::mem::size_of::<usize>() * 8 - hint.leading_zeros() as usize;
+        let mut stack = Vec::with_capacity(cap);
+        self.enumerate().foreach(|(mut i, mut x)| {
+            while (i & 1) != 0 {
+                x = f(stack.pop().unwrap(), x);
+                i >>= 1;
+            }
+            stack.push(x);
+        });
+        stack.into_iter().fold1(f)
+    }
+}
+impl<T:Iterator> IterEx for T {}
+
+macro_rules! def_benchs {
+    ($N:expr,
+     $FUN:ident,
+     $BENCH_NAME:ident,
+     ) => (
+        mod $BENCH_NAME {
+        use super::*;
+
+        #[bench]
+        fn sum(b: &mut Bencher) {
+            let v: Vec<u32> = (0.. $N).collect();
+            b.iter(|| {
+                cloned(&v).$FUN(|x, y| x + y)
+            });
+        }
+
+        #[bench]
+        fn complex_iter(b: &mut Bencher) {
+            let u = (3..).take($N / 2);
+            let v = (5..).take($N / 2);
+            let it = u.chain(v);
+
+            b.iter(|| {
+                it.clone().map(|x| x as f32).$FUN(f32::atan2)
+            });
+        }
+
+        #[bench]
+        fn string_format(b: &mut Bencher) {
+            // This goes quadratic with linear `fold1`, so use a smaller
+            // size to not waste too much time in travis.  The allocations
+            // in here are so expensive anyway that it'll still take
+            // way longer per iteration than the other two benchmarks.
+            let v: Vec<u32> = (0.. ($N/4)).collect();
+            b.iter(|| {
+                cloned(&v).map(|x| x.to_string()).$FUN(|x, y| format!("{} + {}", x, y))
+            });
+        }
+        }
+    )
+}
+
+def_benchs!{
+    10_000,
+    fold1,
+    fold1_10k,
+}
+
+def_benchs!{
+    10_000,
+    tree_fold1,
+    tree_fold1_stack_10k,
+}
+
+def_benchs!{
+    10_000,
+    tree_fold1_vec,
+    tree_fold1_vec_10k,
+}
+
+def_benchs!{
+    100,
+    fold1,
+    fold1_100,
+}
+
+def_benchs!{
+    100,
+    tree_fold1,
+    tree_fold1_stack_100,
+}
+
+def_benchs!{
+    100,
+    tree_fold1_vec,
+    tree_fold1_vec_100,
+}
+
+def_benchs!{
+    8,
+    fold1,
+    fold1_08,
+}
+
+def_benchs!{
+    8,
+    tree_fold1,
+    tree_fold1_stack_08,
+}
+
+def_benchs!{
+    8,
+    tree_fold1_vec,
+    tree_fold1_vec_08,
+}
diff --git a/bors.toml b/bors.toml
deleted file mode 100644
index 2523b2d..0000000
--- a/bors.toml
+++ /dev/null
@@ -1,3 +0,0 @@
-status = [
-"continuous-integration/travis-ci/push"
-]
diff --git a/src/lib.rs b/src/lib.rs
index a66bc38..32b5f71 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1601,6 +1601,114 @@ pub trait Itertools : Iterator {
         self.next().map(move |x| self.fold(x, f))
     }
 
+    /// Accumulate the elements in the iterator in a tree-like manner.
+    ///
+    /// You can think of it as, while there's more than one item, repeatedly
+    /// combining adjacent items.  It does so in bottom-up-merge-sort order,
+    /// however, so that it needs only logarithmic stack space.
+    ///
+    /// This produces a call tree like the following (where the calls under
+    /// an item are done after reading that item):
+    ///
+    /// ```text
+    /// 1 2 3 4 5 6 7
+    /// │ │ │ │ │ │ │
+    /// └─f └─f └─f │
+    ///   │   │   │ │
+    ///   └───f   └─f
+    ///       │     │
+    ///       └─────f
+    /// ```
+    ///
+    /// Which, for non-associative functions, will typically produce a different
+    /// result than the linear call tree used by `fold1`:
+    ///
+    /// ```text
+    /// 1 2 3 4 5 6 7
+    /// │ │ │ │ │ │ │
+    /// └─f─f─f─f─f─f
+    /// ```
+    ///
+    /// If `f` is associative, prefer the normal `fold1` instead.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// // The same tree as above
+    /// let num_strings = (1..8).map(|x| x.to_string());
+    /// assert_eq!(num_strings.tree_fold1(|x, y| format!("f({}, {})", x, y)),
+    ///     Some(String::from("f(f(f(1, 2), f(3, 4)), f(f(5, 6), 7))")));
+    ///
+    /// // Like fold1, an empty iterator produces None
+    /// assert_eq!((0..0).tree_fold1(|x, y| x * y), None);
+    ///
+    /// // tree_fold1 matches fold1 for associative operations...
+    /// assert_eq!((0..10).tree_fold1(|x, y| x + y),
+    ///     (0..10).fold1(|x, y| x + y));
+    /// // ...but not for non-associative ones
+    /// assert!((0..10).tree_fold1(|x, y| x - y)
+    ///     != (0..10).fold1(|x, y| x - y));
+    /// ```
+    // FIXME: If minver changes to >= 1.13, use `assert_ne!` in the doctest.
+    fn tree_fold1<F>(mut self, mut f: F) -> Option<Self::Item>
+        where F: FnMut(Self::Item, Self::Item) -> Self::Item,
+              Self: Sized,
+    {
+        type State<T> = Result<T, Option<T>>;
+
+        fn inner0<T, II, FF>(it: &mut II, f: &mut FF) -> State<T>
+            where
+                II: Iterator<Item = T>,
+                FF: FnMut(T, T) -> T
+        {
+            // This function could be replaced with `it.next().ok_or(None)`,
+            // but half the useful tree_fold1 work is combining adjacent items,
+            // so put that in a form that LLVM is more likely to optimize well.
+
+            let a =
+                if let Some(v) = it.next() { v }
+                else { return Err(None) };
+            let b =
+                if let Some(v) = it.next() { v }
+                else { return Err(Some(a)) };
+            Ok(f(a, b))
+        }
+
+        fn inner<T, II, FF>(stop: usize, it: &mut II, f: &mut FF) -> State<T>
+            where
+                II: Iterator<Item = T>,
+                FF: FnMut(T, T) -> T
+        {
+            let mut x = try!(inner0(it, f));
+            for height in 0..stop {
+                // Try to get another tree the same size with which to combine it,
+                // creating a new tree that's twice as big for next time around.
+                let next =
+                    if height == 0 {
+                        inner0(it, f)
+                    } else {
+                        inner(height, it, f)
+                    };
+                match next {
+                    Ok(y) => x = f(x, y),
+
+                    // If we ran out of items, combine whatever we did manage
+                    // to get.  It's better combined with the current value
+                    // than something in a parent frame, because the tree in
+                    // the parent is always as least as big as this one.
+                    Err(None) => return Err(Some(x)),
+                    Err(Some(y)) => return Err(Some(f(x, y))),
+                }
+            }
+            Ok(x)
+        }
+
+        match inner(usize::max_value(), &mut self, &mut f) {
+            Err(x) => x,
+            _ => unreachable!(),
+        }
+    }
+
     /// An iterator method that applies a function, producing a single, final value.
     ///
     /// `fold_while()` is basically equivalent to `fold()` but with additional support for


Test patch:
diff --git a/tests/quick.rs b/tests/quick.rs
index cb85931..6859758 100644
--- a/tests/quick.rs
+++ b/tests/quick.rs
@@ -986,3 +986,34 @@ quickcheck! {
         TestResult::from_bool(minmax == expected)
     }
 }
+
+quickcheck! {
+    fn tree_fold1_f64(mut a: Vec<f64>) -> TestResult {
+        fn collapse_adjacent<F>(x: Vec<f64>, mut f: F) -> Vec<f64>
+            where F: FnMut(f64, f64) -> f64
+        {
+            let mut out = Vec::new();
+            for i in (0..x.len()).step(2) {
+                if i == x.len()-1 {
+                    out.push(x[i])
+                } else {
+                    out.push(f(x[i], x[i+1]));
+                }
+            }
+            out
+        }
+
+        if a.iter().any(|x| x.is_nan()) {
+            return TestResult::discard();
+        }
+
+        let actual = a.iter().cloned().tree_fold1(f64::atan2);
+
+        while a.len() > 1 {
+            a = collapse_adjacent(a, f64::atan2);
+        }
+        let expected = a.pop();
+
+        TestResult::from_bool(actual == expected)
+    }
+}
diff --git a/tests/test_core.rs b/tests/test_core.rs
index 828e64a..e215255 100644
--- a/tests/test_core.rs
+++ b/tests/test_core.rs
@@ -232,4 +232,9 @@ fn flatten_clone() {
     it::assert_equal(flattened2, &[1,2,3,4,5,6]);
 }
 
-
+#[test]
+fn tree_fold1() {
+    for i in 0..100 {
+        assert_eq!((0..i).tree_fold1(|x, y| x + y), (0..i).fold1(|x, y| x + y));
+    }
+}
diff --git a/tests/test_std.rs b/tests/test_std.rs
index cfd249a..35e0529 100644
--- a/tests/test_std.rs
+++ b/tests/test_std.rs
@@ -719,3 +719,31 @@ fn fold_while() {
     assert_eq!(sum, 15);
 }
 
+#[test]
+fn tree_fold1() {
+    let x = [
+        "",
+        "0",
+        "0 1 x",
+        "0 1 x 2 x",
+        "0 1 x 2 3 x x",
+        "0 1 x 2 3 x x 4 x",
+        "0 1 x 2 3 x x 4 5 x x",
+        "0 1 x 2 3 x x 4 5 x 6 x x",
+        "0 1 x 2 3 x x 4 5 x 6 7 x x x",
+        "0 1 x 2 3 x x 4 5 x 6 7 x x x 8 x",
+        "0 1 x 2 3 x x 4 5 x 6 7 x x x 8 9 x x",
+        "0 1 x 2 3 x x 4 5 x 6 7 x x x 8 9 x 10 x x",
+        "0 1 x 2 3 x x 4 5 x 6 7 x x x 8 9 x 10 11 x x x",
+        "0 1 x 2 3 x x 4 5 x 6 7 x x x 8 9 x 10 11 x x 12 x x",
+        "0 1 x 2 3 x x 4 5 x 6 7 x x x 8 9 x 10 11 x x 12 13 x x x",
+        "0 1 x 2 3 x x 4 5 x 6 7 x x x 8 9 x 10 11 x x 12 13 x 14 x x x",
+        "0 1 x 2 3 x x 4 5 x 6 7 x x x 8 9 x 10 11 x x 12 13 x 14 15 x x x x",
+    ];
+    for (i, &s) in x.iter().enumerate() {
+        let expected = if s == "" { None } else { Some(s.to_string()) };
+        let num_strings = (0..i).map(|x| x.to_string());
+        let actual = num_strings.tree_fold1(|a, b| format!("{} {} x", a, b));
+        assert_eq!(actual, expected);
+    }
+}