rayon/iter/
splitter.rs

1use super::plumbing::*;
2use super::*;
3
4use std::fmt::{self, Debug};
5
6/// The `split` function takes arbitrary data and a closure that knows how to
7/// split it, and turns this into a `ParallelIterator`.
8///
9/// # Examples
10///
11/// As a simple example, Rayon can recursively split ranges of indices
12///
13/// ```
14/// use rayon::iter;
15/// use rayon::prelude::*;
16/// use std::ops::Range;
17///
18///
19/// // We define a range of indices as follows
20/// type Range1D = Range<usize>;
21///
22/// // Splitting it in two can be done like this
23/// fn split_range1(r: Range1D) -> (Range1D, Option<Range1D>) {
24///     // We are mathematically unable to split the range if there is only
25///     // one point inside of it, but we could stop splitting before that.
26///     if r.end - r.start <= 1 { return (r, None); }
27///
28///     // Here, our range is considered large enough to be splittable
29///     let midpoint = r.start + (r.end - r.start) / 2;
30///     (r.start..midpoint, Some(midpoint..r.end))
31/// }
32///
33/// // By using iter::split, Rayon will split the range until it has enough work
34/// // to feed the CPU cores, then give us the resulting sub-ranges
35/// iter::split(0..4096, split_range1).for_each(|sub_range| {
36///     // As our initial range had a power-of-two size, the final sub-ranges
37///     // should have power-of-two sizes too
38///     assert!((sub_range.end - sub_range.start).is_power_of_two());
39/// });
40/// ```
41///
42/// This recursive splitting can be extended to two or three dimensions,
43/// to reproduce a classic "block-wise" parallelization scheme of graphics and
44/// numerical simulations:
45///
46/// ```
47/// # use rayon::iter;
48/// # use rayon::prelude::*;
49/// # use std::ops::Range;
50/// # type Range1D = Range<usize>;
51/// # fn split_range1(r: Range1D) -> (Range1D, Option<Range1D>) {
52/// #     if r.end - r.start <= 1 { return (r, None); }
53/// #     let midpoint = r.start + (r.end - r.start) / 2;
54/// #     (r.start..midpoint, Some(midpoint..r.end))
55/// # }
56/// #
57/// // A two-dimensional range of indices can be built out of two 1D ones
58/// struct Range2D {
59///     // Range of horizontal indices
60///     pub rx: Range1D,
61///
62///     // Range of vertical indices
63///     pub ry: Range1D,
64/// }
65///
66/// // We want to recursively split them by the largest dimension until we have
67/// // enough sub-ranges to feed our mighty multi-core CPU. This function
68/// // carries out one such split.
69/// fn split_range2(r2: Range2D) -> (Range2D, Option<Range2D>) {
70///     // Decide on which axis (horizontal/vertical) the range should be split
71///     let width = r2.rx.end - r2.rx.start;
72///     let height = r2.ry.end - r2.ry.start;
73///     if width >= height {
74///         // This is a wide range, split it on the horizontal axis
75///         let (split_rx, ry) = (split_range1(r2.rx), r2.ry);
76///         let out1 = Range2D {
77///             rx: split_rx.0,
78///             ry: ry.clone(),
79///         };
80///         let out2 = split_rx.1.map(|rx| Range2D { rx, ry });
81///         (out1, out2)
82///     } else {
83///         // This is a tall range, split it on the vertical axis
84///         let (rx, split_ry) = (r2.rx, split_range1(r2.ry));
85///         let out1 = Range2D {
86///             rx: rx.clone(),
87///             ry: split_ry.0,
88///         };
89///         let out2 = split_ry.1.map(|ry| Range2D { rx, ry, });
90///         (out1, out2)
91///     }
92/// }
93///
94/// // Again, rayon can handle the recursive splitting for us
95/// let range = Range2D { rx: 0..800, ry: 0..600 };
96/// iter::split(range, split_range2).for_each(|sub_range| {
97///     // If the sub-ranges were indeed split by the largest dimension, then
98///     // if no dimension was twice larger than the other initially, this
99///     // property will remain true in the final sub-ranges.
100///     let width = sub_range.rx.end - sub_range.rx.start;
101///     let height = sub_range.ry.end - sub_range.ry.start;
102///     assert!((width / 2 <= height) && (height / 2 <= width));
103/// });
104/// ```
105///
106pub fn split<D, S>(data: D, splitter: S) -> Split<D, S>
107where
108    D: Send,
109    S: Fn(D) -> (D, Option<D>) + Sync,
110{
111    Split { data, splitter }
112}
113
114/// `Split` is a parallel iterator using arbitrary data and a splitting function.
115/// This struct is created by the [`split()`] function.
116///
117/// [`split()`]: fn.split.html
118#[derive(Clone)]
119pub struct Split<D, S> {
120    data: D,
121    splitter: S,
122}
123
124impl<D: Debug, S> Debug for Split<D, S> {
125    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
126        f.debug_struct("Split").field("data", &self.data).finish()
127    }
128}
129
130impl<D, S> ParallelIterator for Split<D, S>
131where
132    D: Send,
133    S: Fn(D) -> (D, Option<D>) + Sync + Send,
134{
135    type Item = D;
136
137    fn drive_unindexed<C>(self, consumer: C) -> C::Result
138    where
139        C: UnindexedConsumer<Self::Item>,
140    {
141        let producer = SplitProducer {
142            data: self.data,
143            splitter: &self.splitter,
144        };
145        bridge_unindexed(producer, consumer)
146    }
147}
148
149struct SplitProducer<'a, D, S> {
150    data: D,
151    splitter: &'a S,
152}
153
154impl<'a, D, S> UnindexedProducer for SplitProducer<'a, D, S>
155where
156    D: Send,
157    S: Fn(D) -> (D, Option<D>) + Sync,
158{
159    type Item = D;
160
161    fn split(mut self) -> (Self, Option<Self>) {
162        let splitter = self.splitter;
163        let (left, right) = splitter(self.data);
164        self.data = left;
165        (self, right.map(|data| SplitProducer { data, splitter }))
166    }
167
168    fn fold_with<F>(self, folder: F) -> F
169    where
170        F: Folder<Self::Item>,
171    {
172        folder.consume(self.data)
173    }
174}