Descendants

rayon-demo/src/main.rs

@@ -32,6 +32,8 @@ mod sort;
 #[cfg(test)]
 mod str_split;
 #[cfg(test)]
+mod tree;
+#[cfg(test)]
 mod vec_collect;
 
 #[cfg(test)]

rayon-demo/src/tree.rs

@@ -0,0 +1,79 @@
+//! Some benches for tree walks.
+use rayon::prelude::*;
+
+const SIZE: u64 = 100_000;
+const VAL: u64 = SIZE * (SIZE - 1) / 2;
+
+#[bench]
+fn tree_prefix_collect(b: &mut ::test::Bencher) {
+    let mut vec = None;
+    b.iter(|| {
+        vec = Some(
+            rayon::iter::walk_tree_prefix(0u64..SIZE, |r| {
+                // root is smallest
+                let mid = (r.start + 1 + r.end) / 2;
+                // small indices to the left, large to the right
+                std::iter::once((r.start + 1)..mid)
+                    .chain(std::iter::once(mid..r.end))
+                    .filter(|r| !r.is_empty())
+            })
+            .map(|r| r.start)
+            .collect::<Vec<_>>(),
+        )
+    });
+    assert!(vec.unwrap().into_iter().eq(0..SIZE));
+}
+
+#[bench]
+fn tree_postfix_collect(b: &mut ::test::Bencher) {
+    let mut vec = None;
+    b.iter(|| {
+        vec = Some(
+            rayon::iter::walk_tree_postfix(0u64..SIZE, |r| {
+                // root is largest
+                let mid = (r.start + r.end - 1) / 2;
+                // small indices to the left, large to the right
+                std::iter::once(r.start..mid)
+                    .chain(std::iter::once(mid..(r.end - 1)))
+                    .filter(|r| !r.is_empty())
+            })
+            .map(|r| r.end - 1)
+            .collect::<Vec<_>>(),
+        )
+    });
+    assert!(vec.unwrap().into_iter().eq(0..SIZE));
+}
+
+#[bench]
+fn tree_prefix_sum(b: &mut ::test::Bencher) {
+    b.iter(|| {
+        let s = rayon::iter::walk_tree_prefix(0u64..SIZE, |r| {
+            // root is smallest
+            let mid = (r.start + 1 + r.end) / 2;
+            // small indices to the left, large to the right
+            std::iter::once((r.start + 1)..mid)
+                .chain(std::iter::once(mid..r.end))
+                .filter(|r| !r.is_empty())
+        })
+        .map(|r| r.start)
+        .sum::<u64>();
+        assert_eq!(s, VAL)
+    })
+}
+
+#[bench]
+fn tree_postfix_sum(b: &mut ::test::Bencher) {
+    b.iter(|| {
+        let s = rayon::iter::walk_tree_postfix(0u64..SIZE, |r| {
+            // root is smallest
+            let mid = (r.start + 1 + r.end) / 2;
+            // small indices to the left, large to the right
+            std::iter::once((r.start + 1)..mid)
+                .chain(std::iter::once(mid..r.end))
+                .filter(|r| !r.is_empty())
+        })
+        .map(|r| r.start)
+        .sum::<u64>();
+        assert_eq!(s, VAL)
+    })
+}

src/iter/mod.rs

@@ -155,6 +155,7 @@ mod try_reduce;
 mod try_reduce_with;
 mod unzip;
 mod update;
+mod walk_tree;
 mod while_some;
 mod zip;
 mod zip_eq;
@@ -199,6 +200,9 @@ pub use self::{
     take_any_while::TakeAnyWhile,
     try_fold::{TryFold, TryFoldWith},
     update::Update,
+    walk_tree::{
+        walk_tree, walk_tree_postfix, walk_tree_prefix, WalkTree, WalkTreePostfix, WalkTreePrefix,
+    },
     while_some::WhileSome,
     zip::Zip,
     zip_eq::ZipEq,

src/iter/test.rs

@@ -2240,3 +2240,79 @@ fn check_update() {
 
     assert_eq!(v, vec![vec![1, 0], vec![3, 2, 1, 0]]);
 }
+
+#[test]
+fn walk_tree_prefix() {
+    let v: Vec<u32> = crate::iter::walk_tree_prefix(0u32..100, |r| {
+        // root is smallest
+        let mid = (r.start + 1 + r.end) / 2;
+        // small indices to the left, large to the right
+        std::iter::once((r.start + 1)..mid)
+            .chain(std::iter::once(mid..r.end))
+            .filter(|r| !r.is_empty())
+    })
+    .map(|r| r.start)
+    .collect();
+    assert!(v.into_iter().eq(0..100));
+}
+
+#[test]
+fn walk_tree_postfix() {
+    let v: Vec<_> = crate::iter::walk_tree_postfix(0u64..100, |r| {
+        // root is largest
+        let mid = (r.start + r.end - 1) / 2;
+        // small indices to the left, large to the right
+        std::iter::once(r.start..mid)
+            .chain(std::iter::once(mid..(r.end - 1)))
+            .filter(|r| !r.is_empty())
+    })
+    .map(|r| r.end - 1)
+    .collect();
+    assert!(v.into_iter().eq(0..100));
+}
+
+#[test]
+fn walk_flat_tree_prefix() {
+    let v: Vec<_> =
+        crate::iter::walk_tree_prefix(0, |&e| if e < 99 { Some(e + 1) } else { None }).collect();
+    assert!(v.into_iter().eq(0..100));
+}
+
+#[test]
+fn walk_flat_tree_postfix() {
+    let v: Vec<_> =
+        crate::iter::walk_tree_postfix(99, |&e| if e > 0 { Some(e - 1) } else { None }).collect();
+    assert!(v.into_iter().eq(0..100));
+}
+
+#[test]
+fn walk_tree_prefix_degree5() {
+    let depth = 5;
+    let nodes_number = (1 - 5i32.pow(depth)) / (1 - 5);
+    let nodes = (0..nodes_number).collect::<Vec<_>>();
+    let v: Vec<i32> = crate::iter::walk_tree_prefix(nodes.as_slice(), |&r| {
+        r.split_first()
+            .into_iter()
+            .filter_map(|(_, r)| if r.is_empty() { None } else { Some(r) })
+            .flat_map(|r| r.chunks(r.len() / 5))
+    })
+    .filter_map(|r| r.first().copied())
+    .collect();
+    assert_eq!(v, nodes);
+}
+
+#[test]
+fn walk_tree_postfix_degree5() {
+    let depth = 5;
+    let nodes_number = (1 - 5i32.pow(depth)) / (1 - 5);
+    let nodes = (0..nodes_number).collect::<Vec<_>>();
+    let v: Vec<i32> = crate::iter::walk_tree_postfix(nodes.as_slice(), |&r| {
+        r.split_last()
+            .into_iter()
+            .filter_map(|(_, r)| if r.is_empty() { None } else { Some(r) })
+            .flat_map(|r| r.chunks(r.len() / 5))
+    })
+    .filter_map(|r| r.last().copied())
+    .collect();
+    assert_eq!(v, nodes)
+}