| 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) |
| } |
| } |