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

rayon/iter/
walk_tree.rs
