Revision control
Copy as Markdown
Other Tools
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<S>, // nodes (and subtrees) we have to process
seen: Vec<S>, // nodes which have already been explored
children_of: &'b B, // function generating children
}
impl<S, B, I> UnindexedProducer for WalkTreePrefixProducer<'_, S, B>
where
S: Send,
B: Fn(&S) -> I + Send + Sync,
I: IntoIterator<Item = S>,
I::IntoIter: DoubleEndedIterator,
{
type Item = S;
fn split(mut self) -> (Self, Option<Self>) {
// 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<F>(mut self, mut folder: F) -> F
where
F: Folder<Self::Item>,
{
// 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<S, B> {
initial_state: S,
children_of: B,
}
impl<S, B, I> ParallelIterator for WalkTreePrefix<S, B>
where
S: Send,
B: Fn(&S) -> I + Send + Sync,
I: IntoIterator<Item = S>,
I::IntoIter: DoubleEndedIterator,
{
type Item = S;
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
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::<u32>(), 12);
/// ```
///
/// # Example
///
/// ```
/// use rayon::prelude::*;
/// use rayon::iter::walk_tree_prefix;
///
/// struct Node {
/// content: u32,
/// left: Option<Box<Node>>,
/// right: Option<Box<Node>>,
/// }
///
/// // 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<u32> = 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<S, B, I>(root: S, children_of: B) -> WalkTreePrefix<S, B>
where
S: Send,
B: Fn(&S) -> I + Send + Sync,
I: IntoIterator<Item = S>,
I::IntoIter: DoubleEndedIterator,
{
WalkTreePrefix {
initial_state: root,
children_of,
}
}
// post fix
#[derive(Debug)]
struct WalkTreePostfixProducer<'b, S, B> {
to_explore: Vec<S>, // nodes (and subtrees) we have to process
seen: Vec<S>, // nodes which have already been explored
children_of: &'b B, // function generating children
}
impl<S, B, I> UnindexedProducer for WalkTreePostfixProducer<'_, S, B>
where
S: Send,
B: Fn(&S) -> I + Send + Sync,
I: IntoIterator<Item = S>,
{
type Item = S;
fn split(mut self) -> (Self, Option<Self>) {
// 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<F>(self, mut folder: F) -> F
where
F: Folder<Self::Item>,
{
// 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<F, S, B, I>(children_of: &B, s: S, mut folder: F) -> F
where
F: Folder<S>,
B: Fn(&S) -> I,
I: IntoIterator<Item = S>,
{
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<S, B> {
initial_state: S,
children_of: B,
}
impl<S, B, I> ParallelIterator for WalkTreePostfix<S, B>
where
S: Send,
B: Fn(&S) -> I + Send + Sync,
I: IntoIterator<Item = S>,
{
type Item = S;
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
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<T>(v: &mut Vec<T>) -> Option<Vec<T>> {
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::<u32>(), 12);
/// ```
///
/// # Example
///
/// ```
/// use rayon::prelude::*;
/// use rayon::iter::walk_tree_postfix;
///
/// struct Node {
/// content: u32,
/// left: Option<Box<Node>>,
/// right: Option<Box<Node>>,
/// }
///
/// // 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<u32> = 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<S, B, I>(root: S, children_of: B) -> WalkTreePostfix<S, B>
where
S: Send,
B: Fn(&S) -> I + Send + Sync,
I: IntoIterator<Item = S>,
{
WalkTreePostfix {
initial_state: root,
children_of,
}
}
/// ParallelIterator for arbitrary tree-shaped patterns.
/// Returned by the [`walk_tree()`] function.
#[derive(Debug)]
pub struct WalkTree<S, B>(WalkTreePostfix<S, B>);
/// 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::<u32>(), 12);
/// ```
pub fn walk_tree<S, B, I>(root: S, children_of: B) -> WalkTree<S, B>
where
S: Send,
B: Fn(&S) -> I + Send + Sync,
I: IntoIterator<Item = S>,
I::IntoIter: DoubleEndedIterator,
{
let walker = WalkTreePostfix {
initial_state: root,
children_of,
};
WalkTree(walker)
}
impl<S, B, I> ParallelIterator for WalkTree<S, B>
where
S: Send,
B: Fn(&S) -> I + Send + Sync,
I: IntoIterator<Item = S> + Send,
I::IntoIter: DoubleEndedIterator,
{
type Item = S;
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
self.0.drive_unindexed(consumer)
}
}