rayon/iter/
extend.rs

1use super::noop::NoopConsumer;
2use super::{IntoParallelIterator, ParallelExtend, ParallelIterator};
3
4use std::borrow::Cow;
5use std::collections::LinkedList;
6use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
7use std::collections::{BinaryHeap, VecDeque};
8use std::hash::{BuildHasher, Hash};
9
10/// Performs a generic `par_extend` by collecting to a `LinkedList<Vec<_>>` in
11/// parallel, then extending the collection sequentially.
12fn extend<C, I, F>(collection: &mut C, par_iter: I, reserve: F)
13where
14    I: IntoParallelIterator,
15    F: FnOnce(&mut C, &LinkedList<Vec<I::Item>>),
16    C: Extend<I::Item>,
17{
18    let list = collect(par_iter);
19    reserve(collection, &list);
20    for vec in list {
21        collection.extend(vec);
22    }
23}
24
25pub(super) fn collect<I>(par_iter: I) -> LinkedList<Vec<I::Item>>
26where
27    I: IntoParallelIterator,
28{
29    par_iter
30        .into_par_iter()
31        .fold(Vec::new, vec_push)
32        .map(as_list)
33        .reduce(LinkedList::new, list_append)
34}
35
36fn vec_push<T>(mut vec: Vec<T>, elem: T) -> Vec<T> {
37    vec.push(elem);
38    vec
39}
40
41fn as_list<T>(item: T) -> LinkedList<T> {
42    let mut list = LinkedList::new();
43    list.push_back(item);
44    list
45}
46
47fn list_append<T>(mut list1: LinkedList<T>, mut list2: LinkedList<T>) -> LinkedList<T> {
48    list1.append(&mut list2);
49    list1
50}
51
52/// Computes the total length of a `LinkedList<Vec<_>>`.
53pub(super) fn len<T>(list: &LinkedList<Vec<T>>) -> usize {
54    list.iter().map(Vec::len).sum()
55}
56
57fn no_reserve<C, T>(_: &mut C, _: &LinkedList<Vec<T>>) {}
58
59fn heap_reserve<T: Ord, U>(heap: &mut BinaryHeap<T>, list: &LinkedList<Vec<U>>) {
60    heap.reserve(len(list));
61}
62
63/// Extends a binary heap with items from a parallel iterator.
64impl<T> ParallelExtend<T> for BinaryHeap<T>
65where
66    T: Ord + Send,
67{
68    fn par_extend<I>(&mut self, par_iter: I)
69    where
70        I: IntoParallelIterator<Item = T>,
71    {
72        extend(self, par_iter, heap_reserve);
73    }
74}
75
76/// Extends a binary heap with copied items from a parallel iterator.
77impl<'a, T> ParallelExtend<&'a T> for BinaryHeap<T>
78where
79    T: 'a + Copy + Ord + Send + Sync,
80{
81    fn par_extend<I>(&mut self, par_iter: I)
82    where
83        I: IntoParallelIterator<Item = &'a T>,
84    {
85        extend(self, par_iter, heap_reserve);
86    }
87}
88
89/// Extends a B-tree map with items from a parallel iterator.
90impl<K, V> ParallelExtend<(K, V)> for BTreeMap<K, V>
91where
92    K: Ord + Send,
93    V: Send,
94{
95    fn par_extend<I>(&mut self, par_iter: I)
96    where
97        I: IntoParallelIterator<Item = (K, V)>,
98    {
99        extend(self, par_iter, no_reserve);
100    }
101}
102
103/// Extends a B-tree map with copied items from a parallel iterator.
104impl<'a, K: 'a, V: 'a> ParallelExtend<(&'a K, &'a V)> for BTreeMap<K, V>
105where
106    K: Copy + Ord + Send + Sync,
107    V: Copy + Send + Sync,
108{
109    fn par_extend<I>(&mut self, par_iter: I)
110    where
111        I: IntoParallelIterator<Item = (&'a K, &'a V)>,
112    {
113        extend(self, par_iter, no_reserve);
114    }
115}
116
117/// Extends a B-tree set with items from a parallel iterator.
118impl<T> ParallelExtend<T> for BTreeSet<T>
119where
120    T: Ord + Send,
121{
122    fn par_extend<I>(&mut self, par_iter: I)
123    where
124        I: IntoParallelIterator<Item = T>,
125    {
126        extend(self, par_iter, no_reserve);
127    }
128}
129
130/// Extends a B-tree set with copied items from a parallel iterator.
131impl<'a, T> ParallelExtend<&'a T> for BTreeSet<T>
132where
133    T: 'a + Copy + Ord + Send + Sync,
134{
135    fn par_extend<I>(&mut self, par_iter: I)
136    where
137        I: IntoParallelIterator<Item = &'a T>,
138    {
139        extend(self, par_iter, no_reserve);
140    }
141}
142
143fn map_reserve<K, V, S, U>(map: &mut HashMap<K, V, S>, list: &LinkedList<Vec<U>>)
144where
145    K: Eq + Hash,
146    S: BuildHasher,
147{
148    map.reserve(len(list));
149}
150
151/// Extends a hash map with items from a parallel iterator.
152impl<K, V, S> ParallelExtend<(K, V)> for HashMap<K, V, S>
153where
154    K: Eq + Hash + Send,
155    V: Send,
156    S: BuildHasher + Send,
157{
158    fn par_extend<I>(&mut self, par_iter: I)
159    where
160        I: IntoParallelIterator<Item = (K, V)>,
161    {
162        // See the map_collect benchmarks in rayon-demo for different strategies.
163        extend(self, par_iter, map_reserve);
164    }
165}
166
167/// Extends a hash map with copied items from a parallel iterator.
168impl<'a, K: 'a, V: 'a, S> ParallelExtend<(&'a K, &'a V)> for HashMap<K, V, S>
169where
170    K: Copy + Eq + Hash + Send + Sync,
171    V: Copy + Send + Sync,
172    S: BuildHasher + Send,
173{
174    fn par_extend<I>(&mut self, par_iter: I)
175    where
176        I: IntoParallelIterator<Item = (&'a K, &'a V)>,
177    {
178        extend(self, par_iter, map_reserve);
179    }
180}
181
182fn set_reserve<T, S, U>(set: &mut HashSet<T, S>, list: &LinkedList<Vec<U>>)
183where
184    T: Eq + Hash,
185    S: BuildHasher,
186{
187    set.reserve(len(list));
188}
189
190/// Extends a hash set with items from a parallel iterator.
191impl<T, S> ParallelExtend<T> for HashSet<T, S>
192where
193    T: Eq + Hash + Send,
194    S: BuildHasher + Send,
195{
196    fn par_extend<I>(&mut self, par_iter: I)
197    where
198        I: IntoParallelIterator<Item = T>,
199    {
200        extend(self, par_iter, set_reserve);
201    }
202}
203
204/// Extends a hash set with copied items from a parallel iterator.
205impl<'a, T, S> ParallelExtend<&'a T> for HashSet<T, S>
206where
207    T: 'a + Copy + Eq + Hash + Send + Sync,
208    S: BuildHasher + Send,
209{
210    fn par_extend<I>(&mut self, par_iter: I)
211    where
212        I: IntoParallelIterator<Item = &'a T>,
213    {
214        extend(self, par_iter, set_reserve);
215    }
216}
217
218fn list_push_back<T>(mut list: LinkedList<T>, elem: T) -> LinkedList<T> {
219    list.push_back(elem);
220    list
221}
222
223/// Extends a linked list with items from a parallel iterator.
224impl<T> ParallelExtend<T> for LinkedList<T>
225where
226    T: Send,
227{
228    fn par_extend<I>(&mut self, par_iter: I)
229    where
230        I: IntoParallelIterator<Item = T>,
231    {
232        let mut list = par_iter
233            .into_par_iter()
234            .fold(LinkedList::new, list_push_back)
235            .reduce(LinkedList::new, list_append);
236        self.append(&mut list);
237    }
238}
239
240/// Extends a linked list with copied items from a parallel iterator.
241impl<'a, T> ParallelExtend<&'a T> for LinkedList<T>
242where
243    T: 'a + Copy + Send + Sync,
244{
245    fn par_extend<I>(&mut self, par_iter: I)
246    where
247        I: IntoParallelIterator<Item = &'a T>,
248    {
249        self.par_extend(par_iter.into_par_iter().cloned())
250    }
251}
252
253fn string_push(mut string: String, ch: char) -> String {
254    string.push(ch);
255    string
256}
257
258/// Extends a string with characters from a parallel iterator.
259impl ParallelExtend<char> for String {
260    fn par_extend<I>(&mut self, par_iter: I)
261    where
262        I: IntoParallelIterator<Item = char>,
263    {
264        // This is like `extend`, but `Vec<char>` is less efficient to deal
265        // with than `String`, so instead collect to `LinkedList<String>`.
266        let list: LinkedList<_> = par_iter
267            .into_par_iter()
268            .fold(String::new, string_push)
269            .map(as_list)
270            .reduce(LinkedList::new, list_append);
271
272        self.reserve(list.iter().map(String::len).sum());
273        self.extend(list)
274    }
275}
276
277/// Extends a string with copied characters from a parallel iterator.
278impl<'a> ParallelExtend<&'a char> for String {
279    fn par_extend<I>(&mut self, par_iter: I)
280    where
281        I: IntoParallelIterator<Item = &'a char>,
282    {
283        self.par_extend(par_iter.into_par_iter().cloned())
284    }
285}
286
287fn string_reserve<T: AsRef<str>>(string: &mut String, list: &LinkedList<Vec<T>>) {
288    let len = list.iter().flatten().map(T::as_ref).map(str::len).sum();
289    string.reserve(len);
290}
291
292/// Extends a string with string slices from a parallel iterator.
293impl<'a> ParallelExtend<&'a str> for String {
294    fn par_extend<I>(&mut self, par_iter: I)
295    where
296        I: IntoParallelIterator<Item = &'a str>,
297    {
298        extend(self, par_iter, string_reserve);
299    }
300}
301
302/// Extends a string with strings from a parallel iterator.
303impl ParallelExtend<String> for String {
304    fn par_extend<I>(&mut self, par_iter: I)
305    where
306        I: IntoParallelIterator<Item = String>,
307    {
308        extend(self, par_iter, string_reserve);
309    }
310}
311
312/// Extends a string with string slices from a parallel iterator.
313impl<'a> ParallelExtend<Cow<'a, str>> for String {
314    fn par_extend<I>(&mut self, par_iter: I)
315    where
316        I: IntoParallelIterator<Item = Cow<'a, str>>,
317    {
318        extend(self, par_iter, string_reserve);
319    }
320}
321
322fn deque_reserve<T, U>(deque: &mut VecDeque<T>, list: &LinkedList<Vec<U>>) {
323    deque.reserve(len(list));
324}
325
326/// Extends a deque with items from a parallel iterator.
327impl<T> ParallelExtend<T> for VecDeque<T>
328where
329    T: Send,
330{
331    fn par_extend<I>(&mut self, par_iter: I)
332    where
333        I: IntoParallelIterator<Item = T>,
334    {
335        extend(self, par_iter, deque_reserve);
336    }
337}
338
339/// Extends a deque with copied items from a parallel iterator.
340impl<'a, T> ParallelExtend<&'a T> for VecDeque<T>
341where
342    T: 'a + Copy + Send + Sync,
343{
344    fn par_extend<I>(&mut self, par_iter: I)
345    where
346        I: IntoParallelIterator<Item = &'a T>,
347    {
348        extend(self, par_iter, deque_reserve);
349    }
350}
351
352// See the `collect` module for the `Vec<T>` implementation.
353// impl<T> ParallelExtend<T> for Vec<T>
354
355/// Extends a vector with copied items from a parallel iterator.
356impl<'a, T> ParallelExtend<&'a T> for Vec<T>
357where
358    T: 'a + Copy + Send + Sync,
359{
360    fn par_extend<I>(&mut self, par_iter: I)
361    where
362        I: IntoParallelIterator<Item = &'a T>,
363    {
364        self.par_extend(par_iter.into_par_iter().cloned())
365    }
366}
367
368/// Collapses all unit items from a parallel iterator into one.
369impl ParallelExtend<()> for () {
370    fn par_extend<I>(&mut self, par_iter: I)
371    where
372        I: IntoParallelIterator<Item = ()>,
373    {
374        par_iter.into_par_iter().drive_unindexed(NoopConsumer)
375    }
376}