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<'_, …, …> = 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: ::IntoIter = (children_of)(&s).into_iter();
293 for child: S in children {
294 folder = consume_rec_postfix(children_of, s:child, folder);
295 if folder.full() {
296 return folder;
297 }
298 }
299 folder.consume(item: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<'_, …, …> = 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: usize = v.len() / 2;
339 Some(v.split_off(at: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 = 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}
530