diff --git a/rayon-demo/src/main.rs b/rayon-demo/src/main.rs index 360d80746..bc7f3fac7 100644 --- a/rayon-demo/src/main.rs +++ b/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)] diff --git a/rayon-demo/src/tree.rs b/rayon-demo/src/tree.rs new file mode 100644 index 000000000..1c2c2627e --- /dev/null +++ b/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::>(), + ) + }); + 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::>(), + ) + }); + 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::(); + 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::(); + assert_eq!(s, VAL) + }) +} diff --git a/src/iter/mod.rs b/src/iter/mod.rs index 43897c2ef..8ddc9bab4 100644 --- a/src/iter/mod.rs +++ b/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, diff --git a/src/iter/test.rs b/src/iter/test.rs index 3f9445025..abdd68aa7 100644 --- a/src/iter/test.rs +++ b/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 = 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::>(); + let v: Vec = 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::>(); + let v: Vec = 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) +} diff --git a/src/iter/walk_tree.rs b/src/iter/walk_tree.rs new file mode 100644 index 000000000..f7d6ca86d --- /dev/null +++ b/src/iter/walk_tree.rs @@ -0,0 +1,529 @@ +use crate::iter::plumbing::{bridge_unindexed, Folder, UnindexedConsumer, UnindexedProducer}; +use crate::prelude::*; +use std::iter::once; + +#[derive(Debug)] +struct WalkTreePrefixProducer<'b, S, B> { + to_explore: Vec, // nodes (and subtrees) we have to process + seen: Vec, // nodes which have already been explored + children_of: &'b B, // function generating children +} + +impl UnindexedProducer for WalkTreePrefixProducer<'_, S, B> +where + S: Send, + B: Fn(&S) -> I + Send + Sync, + I: IntoIterator, + I::IntoIter: DoubleEndedIterator, +{ + type Item = S; + + fn split(mut self) -> (Self, Option) { + // explore while front is of size one. + while self.to_explore.len() == 1 { + let front_node = self.to_explore.pop().unwrap(); + self.to_explore + .extend((self.children_of)(&front_node).into_iter().rev()); + self.seen.push(front_node); + } + // now take half of the front. + let right_children = split_vec(&mut self.to_explore); + let right = right_children + .map(|mut c| { + std::mem::swap(&mut c, &mut self.to_explore); + WalkTreePrefixProducer { + to_explore: c, + seen: Vec::new(), + children_of: self.children_of, + } + }) + .or_else(|| { + // we can still try to divide 'seen' + let right_seen = split_vec(&mut self.seen); + right_seen.map(|s| WalkTreePrefixProducer { + to_explore: Default::default(), + seen: s, + children_of: self.children_of, + }) + }); + (self, right) + } + + fn fold_with(mut self, mut folder: F) -> F + where + F: Folder, + { + // start by consuming everything seen + folder = folder.consume_iter(self.seen); + if folder.full() { + return folder; + } + // now do all remaining explorations + while let Some(e) = self.to_explore.pop() { + self.to_explore + .extend((self.children_of)(&e).into_iter().rev()); + folder = folder.consume(e); + if folder.full() { + return folder; + } + } + folder + } +} + +/// ParallelIterator for arbitrary tree-shaped patterns. +/// Returned by the [`walk_tree_prefix()`] function. +#[derive(Debug)] +pub struct WalkTreePrefix { + initial_state: S, + children_of: B, +} + +impl ParallelIterator for WalkTreePrefix +where + S: Send, + B: Fn(&S) -> I + Send + Sync, + I: IntoIterator, + I::IntoIter: DoubleEndedIterator, +{ + type Item = S; + + fn drive_unindexed(self, consumer: C) -> C::Result + where + C: UnindexedConsumer, + { + let producer = WalkTreePrefixProducer { + to_explore: once(self.initial_state).collect(), + seen: Vec::new(), + children_of: &self.children_of, + }; + bridge_unindexed(producer, consumer) + } +} + +/// Create a tree-like prefix parallel iterator from an initial root node. +/// The `children_of` function should take a node and return an iterator over its child nodes. +/// The best parallelization is obtained when the tree is balanced +/// but we should also be able to handle harder cases. +/// +/// # Ordering +/// +/// This function guarantees a prefix ordering. See also [`walk_tree_postfix`], +/// which guarantees a postfix order. +/// If you don't care about ordering, you should use [`walk_tree`], +/// which will use whatever is believed to be fastest. +/// For example a perfect binary tree of 7 nodes will reduced in the following order: +/// +/// ```text +/// a +/// / \ +/// / \ +/// b c +/// / \ / \ +/// d e f g +/// +/// reduced as a,b,d,e,c,f,g +/// +/// ``` +/// +/// # Example +/// +/// ```text +/// 4 +/// / \ +/// / \ +/// 2 3 +/// / \ +/// 1 2 +/// ``` +/// +/// ``` +/// use rayon::iter::walk_tree_prefix; +/// use rayon::prelude::*; +/// +/// let par_iter = walk_tree_prefix(4, |&e| { +/// if e <= 2 { +/// Vec::new() +/// } else { +/// vec![e / 2, e / 2 + 1] +/// } +/// }); +/// assert_eq!(par_iter.sum::(), 12); +/// ``` +/// +/// # Example +/// +/// ``` +/// use rayon::prelude::*; +/// use rayon::iter::walk_tree_prefix; +/// +/// struct Node { +/// content: u32, +/// left: Option>, +/// right: Option>, +/// } +/// +/// // Here we loop on the following tree: +/// // +/// // 10 +/// // / \ +/// // / \ +/// // 3 14 +/// // \ +/// // \ +/// // 18 +/// +/// let root = Node { +/// content: 10, +/// left: Some(Box::new(Node { +/// content: 3, +/// left: None, +/// right: None, +/// })), +/// right: Some(Box::new(Node { +/// content: 14, +/// left: None, +/// right: Some(Box::new(Node { +/// content: 18, +/// left: None, +/// right: None, +/// })), +/// })), +/// }; +/// +/// let mut v: Vec = walk_tree_prefix(&root, |r| { +/// r.left +/// .as_ref() +/// .into_iter() +/// .chain(r.right.as_ref()) +/// .map(|n| &**n) +/// }) +/// .map(|node| node.content) +/// .collect(); +/// assert_eq!(v, vec![10, 3, 14, 18]); +/// ``` +/// +pub fn walk_tree_prefix(root: S, children_of: B) -> WalkTreePrefix +where + S: Send, + B: Fn(&S) -> I + Send + Sync, + I: IntoIterator, + I::IntoIter: DoubleEndedIterator, +{ + WalkTreePrefix { + initial_state: root, + children_of, + } +} + +// post fix + +#[derive(Debug)] +struct WalkTreePostfixProducer<'b, S, B> { + to_explore: Vec, // nodes (and subtrees) we have to process + seen: Vec, // nodes which have already been explored + children_of: &'b B, // function generating children +} + +impl UnindexedProducer for WalkTreePostfixProducer<'_, S, B> +where + S: Send, + B: Fn(&S) -> I + Send + Sync, + I: IntoIterator, +{ + type Item = S; + + fn split(mut self) -> (Self, Option) { + // explore while front is of size one. + while self.to_explore.len() == 1 { + let front_node = self.to_explore.pop().unwrap(); + self.to_explore + .extend((self.children_of)(&front_node).into_iter()); + self.seen.push(front_node); + } + // now take half of the front. + let right_children = split_vec(&mut self.to_explore); + let right = right_children + .map(|c| { + let right_seen = std::mem::take(&mut self.seen); // postfix -> upper nodes are processed last + WalkTreePostfixProducer { + to_explore: c, + seen: right_seen, + children_of: self.children_of, + } + }) + .or_else(|| { + // we can still try to divide 'seen' + let right_seen = split_vec(&mut self.seen); + right_seen.map(|mut s| { + std::mem::swap(&mut self.seen, &mut s); + WalkTreePostfixProducer { + to_explore: Default::default(), + seen: s, + children_of: self.children_of, + } + }) + }); + (self, right) + } + + fn fold_with(self, mut folder: F) -> F + where + F: Folder, + { + // now do all remaining explorations + for e in self.to_explore { + folder = consume_rec_postfix(&self.children_of, e, folder); + if folder.full() { + return folder; + } + } + // end by consuming everything seen + folder.consume_iter(self.seen.into_iter().rev()) + } +} + +fn consume_rec_postfix(children_of: &B, s: S, mut folder: F) -> F +where + F: Folder, + B: Fn(&S) -> I, + I: IntoIterator, +{ + let children = (children_of)(&s).into_iter(); + for child in children { + folder = consume_rec_postfix(children_of, child, folder); + if folder.full() { + return folder; + } + } + folder.consume(s) +} + +/// ParallelIterator for arbitrary tree-shaped patterns. +/// Returned by the [`walk_tree_postfix()`] function. +#[derive(Debug)] +pub struct WalkTreePostfix { + initial_state: S, + children_of: B, +} + +impl ParallelIterator for WalkTreePostfix +where + S: Send, + B: Fn(&S) -> I + Send + Sync, + I: IntoIterator, +{ + type Item = S; + + fn drive_unindexed(self, consumer: C) -> C::Result + where + C: UnindexedConsumer, + { + let producer = WalkTreePostfixProducer { + to_explore: once(self.initial_state).collect(), + seen: Vec::new(), + children_of: &self.children_of, + }; + bridge_unindexed(producer, consumer) + } +} + +/// Divide given vector in two equally sized vectors. +/// Return `None` if initial size is <=1. +/// We return the first half and keep the last half in `v`. +fn split_vec(v: &mut Vec) -> Option> { + if v.len() <= 1 { + None + } else { + let n = v.len() / 2; + Some(v.split_off(n)) + } +} + +/// Create a tree like postfix parallel iterator from an initial root node. +/// The `children_of` function should take a node and iterate on all of its child nodes. +/// The best parallelization is obtained when the tree is balanced +/// but we should also be able to handle harder cases. +/// +/// # Ordering +/// +/// This function guarantees a postfix ordering. See also [`walk_tree_prefix`] which guarantees a +/// prefix order. If you don't care about ordering, you should use [`walk_tree`], which will use +/// whatever is believed to be fastest. +/// +/// Between siblings, children are reduced in order -- that is first children are reduced first. +/// +/// For example a perfect binary tree of 7 nodes will reduced in the following order: +/// +/// ```text +/// a +/// / \ +/// / \ +/// b c +/// / \ / \ +/// d e f g +/// +/// reduced as d,e,b,f,g,c,a +/// +/// ``` +/// +/// # Example +/// +/// ```text +/// 4 +/// / \ +/// / \ +/// 2 3 +/// / \ +/// 1 2 +/// ``` +/// +/// ``` +/// use rayon::iter::walk_tree_postfix; +/// use rayon::prelude::*; +/// +/// let par_iter = walk_tree_postfix(4, |&e| { +/// if e <= 2 { +/// Vec::new() +/// } else { +/// vec![e / 2, e / 2 + 1] +/// } +/// }); +/// assert_eq!(par_iter.sum::(), 12); +/// ``` +/// +/// # Example +/// +/// ``` +/// use rayon::prelude::*; +/// use rayon::iter::walk_tree_postfix; +/// +/// struct Node { +/// content: u32, +/// left: Option>, +/// right: Option>, +/// } +/// +/// // Here we loop on the following tree: +/// // +/// // 10 +/// // / \ +/// // / \ +/// // 3 14 +/// // \ +/// // \ +/// // 18 +/// +/// let root = Node { +/// content: 10, +/// left: Some(Box::new(Node { +/// content: 3, +/// left: None, +/// right: None, +/// })), +/// right: Some(Box::new(Node { +/// content: 14, +/// left: None, +/// right: Some(Box::new(Node { +/// content: 18, +/// left: None, +/// right: None, +/// })), +/// })), +/// }; +/// +/// let mut v: Vec = walk_tree_postfix(&root, |r| { +/// r.left +/// .as_ref() +/// .into_iter() +/// .chain(r.right.as_ref()) +/// .map(|n| &**n) +/// }) +/// .map(|node| node.content) +/// .collect(); +/// assert_eq!(v, vec![3, 18, 14, 10]); +/// ``` +/// +pub fn walk_tree_postfix(root: S, children_of: B) -> WalkTreePostfix +where + S: Send, + B: Fn(&S) -> I + Send + Sync, + I: IntoIterator, +{ + WalkTreePostfix { + initial_state: root, + children_of, + } +} + +/// ParallelIterator for arbitrary tree-shaped patterns. +/// Returned by the [`walk_tree()`] function. +#[derive(Debug)] +pub struct WalkTree(WalkTreePostfix); + +/// Create a tree like parallel iterator from an initial root node. +/// The `children_of` function should take a node and iterate on all of its child nodes. +/// The best parallelization is obtained when the tree is balanced +/// but we should also be able to handle harder cases. +/// +/// # Ordering +/// +/// This function does not guarantee any ordering but will +/// use whatever algorithm is thought to achieve the fastest traversal. +/// See also [`walk_tree_prefix`] which guarantees a +/// prefix order and [`walk_tree_postfix`] which guarantees a postfix order. +/// +/// # Example +/// +/// ```text +/// 4 +/// / \ +/// / \ +/// 2 3 +/// / \ +/// 1 2 +/// ``` +/// +/// ``` +/// use rayon::iter::walk_tree; +/// use rayon::prelude::*; +/// +/// let par_iter = walk_tree(4, |&e| { +/// if e <= 2 { +/// Vec::new() +/// } else { +/// vec![e / 2, e / 2 + 1] +/// } +/// }); +/// assert_eq!(par_iter.sum::(), 12); +/// ``` +pub fn walk_tree(root: S, children_of: B) -> WalkTree +where + S: Send, + B: Fn(&S) -> I + Send + Sync, + I: IntoIterator, + I::IntoIter: DoubleEndedIterator, +{ + let walker = WalkTreePostfix { + initial_state: root, + children_of, + }; + WalkTree(walker) +} + +impl ParallelIterator for WalkTree +where + S: Send, + B: Fn(&S) -> I + Send + Sync, + I: IntoIterator + Send, + I::IntoIter: DoubleEndedIterator, +{ + type Item = S; + + fn drive_unindexed(self, consumer: C) -> C::Result + where + C: UnindexedConsumer, + { + self.0.drive_unindexed(consumer) + } +}