rayon/iter/
walk_tree.rs

1use crate::iter::plumbing::{bridge_unindexed, Folder, UnindexedConsumer, UnindexedProducer};
2use crate::prelude::*;
3use std::iter::once;
4
5#[derive(Debug)]
6struct WalkTreePrefixProducer<'b, S, B> {
7    to_explore: Vec<S>, // nodes (and subtrees) we have to process
8    seen: Vec<S>,       // nodes which have already been explored
9    children_of: &'b B, // function generating children
10}
11
12impl<S, B, I> UnindexedProducer for WalkTreePrefixProducer<'_, S, B>
13where
14    S: Send,
15    B: Fn(&S) -> I + Send + Sync,
16    I: IntoIterator<Item = S>,
17    I::IntoIter: DoubleEndedIterator,
18{
19    type Item = S;
20
21    fn split(mut self) -> (Self, Option<Self>) {
22        // explore while front is of size one.
23        while self.to_explore.len() == 1 {
24            let front_node = self.to_explore.pop().unwrap();
25            self.to_explore
26                .extend((self.children_of)(&front_node).into_iter().rev());
27            self.seen.push(front_node);
28        }
29        // now take half of the front.
30        let right_children = split_vec(&mut self.to_explore);
31        let right = right_children
32            .map(|mut c| {
33                std::mem::swap(&mut c, &mut self.to_explore);
34                WalkTreePrefixProducer {
35                    to_explore: c,
36                    seen: Vec::new(),
37                    children_of: self.children_of,
38                }
39            })
40            .or_else(|| {
41                // we can still try to divide 'seen'
42                let right_seen = split_vec(&mut self.seen);
43                right_seen.map(|s| WalkTreePrefixProducer {
44                    to_explore: Default::default(),
45                    seen: s,
46                    children_of: self.children_of,
47                })
48            });
49        (self, right)
50    }
51
52    fn fold_with<F>(mut self, mut folder: F) -> F
53    where
54        F: Folder<Self::Item>,
55    {
56        // start by consuming everything seen
57        folder = folder.consume_iter(self.seen);
58        if folder.full() {
59            return folder;
60        }
61        // now do all remaining explorations
62        while let Some(e) = self.to_explore.pop() {
63            self.to_explore
64                .extend((self.children_of)(&e).into_iter().rev());
65            folder = folder.consume(e);
66            if folder.full() {
67                return folder;
68            }
69        }
70        folder
71    }
72}
73
74/// ParallelIterator for arbitrary tree-shaped patterns.
75/// Returned by the [`walk_tree_prefix()`] function.
76#[derive(Debug)]
77pub struct WalkTreePrefix<S, B> {
78    initial_state: S,
79    children_of: B,
80}
81
82impl<S, B, I> ParallelIterator for WalkTreePrefix<S, B>
83where
84    S: Send,
85    B: Fn(&S) -> I + Send + Sync,
86    I: IntoIterator<Item = S>,
87    I::IntoIter: DoubleEndedIterator,
88{
89    type Item = S;
90
91    fn drive_unindexed<C>(self, consumer: C) -> C::Result
92    where
93        C: UnindexedConsumer<Self::Item>,
94    {
95        let producer = WalkTreePrefixProducer {
96            to_explore: once(self.initial_state).collect(),
97            seen: Vec::new(),
98            children_of: &self.children_of,
99        };
100        bridge_unindexed(producer, consumer)
101    }
102}
103
104/// Create a tree-like prefix parallel iterator from an initial root node.
105/// The `children_of` function should take a node and return an iterator over its child nodes.
106/// The best parallelization is obtained when the tree is balanced
107/// but we should also be able to handle harder cases.
108///
109/// # Ordering
110///
111/// This function guarantees a prefix ordering. See also [`walk_tree_postfix`],
112/// which guarantees a postfix order.
113/// If you don't care about ordering, you should use [`walk_tree`],
114/// which will use whatever is believed to be fastest.
115/// For example a perfect binary tree of 7 nodes will reduced in the following order:
116///
117/// ```text
118///      a
119///     / \
120///    /   \
121///   b     c
122///  / \   / \
123/// d   e f   g
124///
125/// reduced as a,b,d,e,c,f,g
126///
127/// ```
128///
129/// # Example
130///
131/// ```text
132///      4
133///     / \
134///    /   \
135///   2     3
136///        / \
137///       1   2
138/// ```
139///
140/// ```
141/// use rayon::iter::walk_tree_prefix;
142/// use rayon::prelude::*;
143///
144/// let par_iter = walk_tree_prefix(4, |&e| {
145///     if e <= 2 {
146///         Vec::new()
147///     } else {
148///         vec![e / 2, e / 2 + 1]
149///     }
150/// });
151/// assert_eq!(par_iter.sum::<u32>(), 12);
152/// ```
153///
154/// # Example
155///
156/// ```
157/// use rayon::prelude::*;
158/// use rayon::iter::walk_tree_prefix;
159///
160/// struct Node {
161///     content: u32,
162///     left: Option<Box<Node>>,
163///     right: Option<Box<Node>>,
164/// }
165///
166/// // Here we loop on the following tree:
167/// //
168/// //       10
169/// //      /  \
170/// //     /    \
171/// //    3     14
172/// //            \
173/// //             \
174/// //              18
175///
176/// let root = Node {
177///     content: 10,
178///     left: Some(Box::new(Node {
179///         content: 3,
180///         left: None,
181///         right: None,
182///     })),
183///     right: Some(Box::new(Node {
184///         content: 14,
185///         left: None,
186///         right: Some(Box::new(Node {
187///             content: 18,
188///             left: None,
189///             right: None,
190///         })),
191///     })),
192/// };
193///
194/// let mut v: Vec<u32> = walk_tree_prefix(&root, |r| {
195///     r.left
196///         .as_ref()
197///         .into_iter()
198///         .chain(r.right.as_ref())
199///         .map(|n| &**n)
200/// })
201/// .map(|node| node.content)
202/// .collect();
203/// assert_eq!(v, vec![10, 3, 14, 18]);
204/// ```
205///
206pub fn walk_tree_prefix<S, B, I>(root: S, children_of: B) -> WalkTreePrefix<S, B>
207where
208    S: Send,
209    B: Fn(&S) -> I + Send + Sync,
210    I: IntoIterator<Item = S>,
211    I::IntoIter: DoubleEndedIterator,
212{
213    WalkTreePrefix {
214        initial_state: root,
215        children_of,
216    }
217}
218
219// post fix
220
221#[derive(Debug)]
222struct WalkTreePostfixProducer<'b, S, B> {
223    to_explore: Vec<S>, // nodes (and subtrees) we have to process
224    seen: Vec<S>,       // nodes which have already been explored
225    children_of: &'b B, // function generating children
226}
227
228impl<S, B, I> UnindexedProducer for WalkTreePostfixProducer<'_, S, B>
229where
230    S: Send,
231    B: Fn(&S) -> I + Send + Sync,
232    I: IntoIterator<Item = S>,
233{
234    type Item = S;
235
236    fn split(mut self) -> (Self, Option<Self>) {
237        // explore while front is of size one.
238        while self.to_explore.len() == 1 {
239            let front_node = self.to_explore.pop().unwrap();
240            self.to_explore
241                .extend((self.children_of)(&front_node).into_iter());
242            self.seen.push(front_node);
243        }
244        // now take half of the front.
245        let right_children = split_vec(&mut self.to_explore);
246        let right = right_children
247            .map(|c| {
248                let right_seen = std::mem::take(&mut self.seen); // postfix -> upper nodes are processed last
249                WalkTreePostfixProducer {
250                    to_explore: c,
251                    seen: right_seen,
252                    children_of: self.children_of,
253                }
254            })
255            .or_else(|| {
256                // we can still try to divide 'seen'
257                let right_seen = split_vec(&mut self.seen);
258                right_seen.map(|mut s| {
259                    std::mem::swap(&mut self.seen, &mut s);
260                    WalkTreePostfixProducer {
261                        to_explore: Default::default(),
262                        seen: s,
263                        children_of: self.children_of,
264                    }
265                })
266            });
267        (self, right)
268    }
269
270    fn fold_with<F>(self, mut folder: F) -> F
271    where
272        F: Folder<Self::Item>,
273    {
274        // now do all remaining explorations
275        for e in self.to_explore {
276            folder = consume_rec_postfix(&self.children_of, e, folder);
277            if folder.full() {
278                return folder;
279            }
280        }
281        // end by consuming everything seen
282        folder.consume_iter(self.seen.into_iter().rev())
283    }
284}
285
286fn consume_rec_postfix<F, S, B, I>(children_of: &B, s: S, mut folder: F) -> F
287where
288    F: Folder<S>,
289    B: Fn(&S) -> I,
290    I: IntoIterator<Item = S>,
291{
292    let children = (children_of)(&s).into_iter();
293    for child in children {
294        folder = consume_rec_postfix(children_of, child, folder);
295        if folder.full() {
296            return folder;
297        }
298    }
299    folder.consume(s)
300}
301
302/// ParallelIterator for arbitrary tree-shaped patterns.
303/// Returned by the [`walk_tree_postfix()`] function.
304#[derive(Debug)]
305pub struct WalkTreePostfix<S, B> {
306    initial_state: S,
307    children_of: B,
308}
309
310impl<S, B, I> ParallelIterator for WalkTreePostfix<S, B>
311where
312    S: Send,
313    B: Fn(&S) -> I + Send + Sync,
314    I: IntoIterator<Item = S>,
315{
316    type Item = S;
317
318    fn drive_unindexed<C>(self, consumer: C) -> C::Result
319    where
320        C: UnindexedConsumer<Self::Item>,
321    {
322        let producer = WalkTreePostfixProducer {
323            to_explore: once(self.initial_state).collect(),
324            seen: Vec::new(),
325            children_of: &self.children_of,
326        };
327        bridge_unindexed(producer, consumer)
328    }
329}
330
331/// Divide given vector in two equally sized vectors.
332/// Return `None` if initial size is <=1.
333/// We return the first half and keep the last half in `v`.
334fn split_vec<T>(v: &mut Vec<T>) -> Option<Vec<T>> {
335    if v.len() <= 1 {
336        None
337    } else {
338        let n = v.len() / 2;
339        Some(v.split_off(n))
340    }
341}
342
343/// Create a tree like postfix parallel iterator from an initial root node.
344/// The `children_of` function should take a node and iterate on all of its child nodes.
345/// The best parallelization is obtained when the tree is balanced
346/// but we should also be able to handle harder cases.
347///
348/// # Ordering
349///
350/// This function guarantees a postfix ordering. See also [`walk_tree_prefix`] which guarantees a
351/// prefix order. If you don't care about ordering, you should use [`walk_tree`], which will use
352/// whatever is believed to be fastest.
353///
354/// Between siblings, children are reduced in order -- that is first children are reduced first.
355///
356/// For example a perfect binary tree of 7 nodes will reduced in the following order:
357///
358/// ```text
359///      a
360///     / \
361///    /   \
362///   b     c
363///  / \   / \
364/// d   e f   g
365///
366/// reduced as d,e,b,f,g,c,a
367///
368/// ```
369///
370/// # Example
371///
372/// ```text
373///      4
374///     / \
375///    /   \
376///   2     3
377///        / \
378///       1   2
379/// ```
380///
381/// ```
382/// use rayon::iter::walk_tree_postfix;
383/// use rayon::prelude::*;
384///
385/// let par_iter = walk_tree_postfix(4, |&e| {
386///     if e <= 2 {
387///         Vec::new()
388///     } else {
389///         vec![e / 2, e / 2 + 1]
390///     }
391/// });
392/// assert_eq!(par_iter.sum::<u32>(), 12);
393/// ```
394///
395/// # Example
396///
397/// ```
398/// use rayon::prelude::*;
399/// use rayon::iter::walk_tree_postfix;
400///
401/// struct Node {
402///     content: u32,
403///     left: Option<Box<Node>>,
404///     right: Option<Box<Node>>,
405/// }
406///
407/// // Here we loop on the following tree:
408/// //
409/// //       10
410/// //      /  \
411/// //     /    \
412/// //    3     14
413/// //            \
414/// //             \
415/// //              18
416///
417/// let root = Node {
418///     content: 10,
419///     left: Some(Box::new(Node {
420///         content: 3,
421///         left: None,
422///         right: None,
423///     })),
424///     right: Some(Box::new(Node {
425///         content: 14,
426///         left: None,
427///         right: Some(Box::new(Node {
428///             content: 18,
429///             left: None,
430///             right: None,
431///         })),
432///     })),
433/// };
434///
435/// let mut v: Vec<u32> = walk_tree_postfix(&root, |r| {
436///     r.left
437///         .as_ref()
438///         .into_iter()
439///         .chain(r.right.as_ref())
440///         .map(|n| &**n)
441/// })
442/// .map(|node| node.content)
443/// .collect();
444/// assert_eq!(v, vec![3, 18, 14, 10]);
445/// ```
446///
447pub fn walk_tree_postfix<S, B, I>(root: S, children_of: B) -> WalkTreePostfix<S, B>
448where
449    S: Send,
450    B: Fn(&S) -> I + Send + Sync,
451    I: IntoIterator<Item = S>,
452{
453    WalkTreePostfix {
454        initial_state: root,
455        children_of,
456    }
457}
458
459/// ParallelIterator for arbitrary tree-shaped patterns.
460/// Returned by the [`walk_tree()`] function.
461#[derive(Debug)]
462pub struct WalkTree<S, B>(WalkTreePostfix<S, B>);
463
464/// Create a tree like parallel iterator from an initial root node.
465/// The `children_of` function should take a node and iterate on all of its child nodes.
466/// The best parallelization is obtained when the tree is balanced
467/// but we should also be able to handle harder cases.
468///
469/// # Ordering
470///
471/// This function does not guarantee any ordering but will
472/// use whatever algorithm is thought to achieve the fastest traversal.
473/// See also [`walk_tree_prefix`] which guarantees a
474/// prefix order and [`walk_tree_postfix`] which guarantees a postfix order.
475///
476/// # Example
477///
478/// ```text
479///      4
480///     / \
481///    /   \
482///   2     3
483///        / \
484///       1   2
485/// ```
486///
487/// ```
488/// use rayon::iter::walk_tree;
489/// use rayon::prelude::*;
490///
491/// let par_iter = walk_tree(4, |&e| {
492///     if e <= 2 {
493///         Vec::new()
494///     } else {
495///         vec![e / 2, e / 2 + 1]
496///     }
497/// });
498/// assert_eq!(par_iter.sum::<u32>(), 12);
499/// ```
500pub fn walk_tree<S, B, I>(root: S, children_of: B) -> WalkTree<S, B>
501where
502    S: Send,
503    B: Fn(&S) -> I + Send + Sync,
504    I: IntoIterator<Item = S>,
505    I::IntoIter: DoubleEndedIterator,
506{
507    let walker = WalkTreePostfix {
508        initial_state: root,
509        children_of,
510    };
511    WalkTree(walker)
512}
513
514impl<S, B, I> ParallelIterator for WalkTree<S, B>
515where
516    S: Send,
517    B: Fn(&S) -> I + Send + Sync,
518    I: IntoIterator<Item = S> + Send,
519    I::IntoIter: DoubleEndedIterator,
520{
521    type Item = S;
522
523    fn drive_unindexed<C>(self, consumer: C) -> C::Result
524    where
525        C: UnindexedConsumer<Self::Item>,
526    {
527        self.0.drive_unindexed(consumer)
528    }
529}
rayon/iter/walk_tree.rs

rayon/iter/
walk_tree.rs
