rayon/iter/
map_with.rs

1use super::plumbing::*;
2use super::*;
3
4use std::fmt::{self, Debug};
5
6/// `MapWith` is an iterator that transforms the elements of an underlying iterator.
7///
8/// This struct is created by the [`map_with()`] method on [`ParallelIterator`]
9///
10/// [`map_with()`]: ParallelIterator::map_with()
11#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
12#[derive(Clone)]
13pub struct MapWith<I, T, F> {
14    base: I,
15    item: T,
16    map_op: F,
17}
18
19impl<I: Debug, T: Debug, F> Debug for MapWith<I, T, F> {
20    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
21        f.debug_struct("MapWith")
22            .field("base", &self.base)
23            .field("item", &self.item)
24            .finish()
25    }
26}
27
28impl<I, T, F> MapWith<I, T, F> {
29    /// Creates a new `MapWith` iterator.
30    pub(super) fn new(base: I, item: T, map_op: F) -> Self {
31        MapWith { base, item, map_op }
32    }
33}
34
35impl<I, T, F, R> ParallelIterator for MapWith<I, T, F>
36where
37    I: ParallelIterator,
38    T: Send + Clone,
39    F: Fn(&mut T, I::Item) -> R + Sync + Send,
40    R: Send,
41{
42    type Item = R;
43
44    fn drive_unindexed<C>(self, consumer: C) -> C::Result
45    where
46        C: UnindexedConsumer<Self::Item>,
47    {
48        let consumer1 = MapWithConsumer::new(consumer, self.item, &self.map_op);
49        self.base.drive_unindexed(consumer1)
50    }
51
52    fn opt_len(&self) -> Option<usize> {
53        self.base.opt_len()
54    }
55}
56
57impl<I, T, F, R> IndexedParallelIterator for MapWith<I, T, F>
58where
59    I: IndexedParallelIterator,
60    T: Send + Clone,
61    F: Fn(&mut T, I::Item) -> R + Sync + Send,
62    R: Send,
63{
64    fn drive<C>(self, consumer: C) -> C::Result
65    where
66        C: Consumer<Self::Item>,
67    {
68        let consumer1 = MapWithConsumer::new(consumer, self.item, &self.map_op);
69        self.base.drive(consumer1)
70    }
71
72    fn len(&self) -> usize {
73        self.base.len()
74    }
75
76    fn with_producer<CB>(self, callback: CB) -> CB::Output
77    where
78        CB: ProducerCallback<Self::Item>,
79    {
80        return self.base.with_producer(Callback {
81            callback,
82            item: self.item,
83            map_op: self.map_op,
84        });
85
86        struct Callback<CB, U, F> {
87            callback: CB,
88            item: U,
89            map_op: F,
90        }
91
92        impl<T, U, F, R, CB> ProducerCallback<T> for Callback<CB, U, F>
93        where
94            CB: ProducerCallback<R>,
95            U: Send + Clone,
96            F: Fn(&mut U, T) -> R + Sync,
97            R: Send,
98        {
99            type Output = CB::Output;
100
101            fn callback<P>(self, base: P) -> CB::Output
102            where
103                P: Producer<Item = T>,
104            {
105                let producer = MapWithProducer {
106                    base,
107                    item: self.item,
108                    map_op: &self.map_op,
109                };
110                self.callback.callback(producer)
111            }
112        }
113    }
114}
115
116// ////////////////////////////////////////////////////////////////////////
117
118struct MapWithProducer<'f, P, U, F> {
119    base: P,
120    item: U,
121    map_op: &'f F,
122}
123
124impl<'f, P, U, F, R> Producer for MapWithProducer<'f, P, U, F>
125where
126    P: Producer,
127    U: Send + Clone,
128    F: Fn(&mut U, P::Item) -> R + Sync,
129    R: Send,
130{
131    type Item = R;
132    type IntoIter = MapWithIter<'f, P::IntoIter, U, F>;
133
134    fn into_iter(self) -> Self::IntoIter {
135        MapWithIter {
136            base: self.base.into_iter(),
137            item: self.item,
138            map_op: self.map_op,
139        }
140    }
141
142    fn min_len(&self) -> usize {
143        self.base.min_len()
144    }
145    fn max_len(&self) -> usize {
146        self.base.max_len()
147    }
148
149    fn split_at(self, index: usize) -> (Self, Self) {
150        let (left, right) = self.base.split_at(index);
151        (
152            MapWithProducer {
153                base: left,
154                item: self.item.clone(),
155                map_op: self.map_op,
156            },
157            MapWithProducer {
158                base: right,
159                item: self.item,
160                map_op: self.map_op,
161            },
162        )
163    }
164
165    fn fold_with<G>(self, folder: G) -> G
166    where
167        G: Folder<Self::Item>,
168    {
169        let folder1 = MapWithFolder {
170            base: folder,
171            item: self.item,
172            map_op: self.map_op,
173        };
174        self.base.fold_with(folder1).base
175    }
176}
177
178struct MapWithIter<'f, I, U, F> {
179    base: I,
180    item: U,
181    map_op: &'f F,
182}
183
184impl<'f, I, U, F, R> Iterator for MapWithIter<'f, I, U, F>
185where
186    I: Iterator,
187    F: Fn(&mut U, I::Item) -> R + Sync,
188    R: Send,
189{
190    type Item = R;
191
192    fn next(&mut self) -> Option<R> {
193        let item = self.base.next()?;
194        Some((self.map_op)(&mut self.item, item))
195    }
196
197    fn size_hint(&self) -> (usize, Option<usize>) {
198        self.base.size_hint()
199    }
200}
201
202impl<'f, I, U, F, R> DoubleEndedIterator for MapWithIter<'f, I, U, F>
203where
204    I: DoubleEndedIterator,
205    F: Fn(&mut U, I::Item) -> R + Sync,
206    R: Send,
207{
208    fn next_back(&mut self) -> Option<R> {
209        let item = self.base.next_back()?;
210        Some((self.map_op)(&mut self.item, item))
211    }
212}
213
214impl<'f, I, U, F, R> ExactSizeIterator for MapWithIter<'f, I, U, F>
215where
216    I: ExactSizeIterator,
217    F: Fn(&mut U, I::Item) -> R + Sync,
218    R: Send,
219{
220}
221
222// ////////////////////////////////////////////////////////////////////////
223// Consumer implementation
224
225struct MapWithConsumer<'f, C, U, F> {
226    base: C,
227    item: U,
228    map_op: &'f F,
229}
230
231impl<'f, C, U, F> MapWithConsumer<'f, C, U, F> {
232    fn new(base: C, item: U, map_op: &'f F) -> Self {
233        MapWithConsumer { base, item, map_op }
234    }
235}
236
237impl<'f, T, U, R, C, F> Consumer<T> for MapWithConsumer<'f, C, U, F>
238where
239    C: Consumer<R>,
240    U: Send + Clone,
241    F: Fn(&mut U, T) -> R + Sync,
242    R: Send,
243{
244    type Folder = MapWithFolder<'f, C::Folder, U, F>;
245    type Reducer = C::Reducer;
246    type Result = C::Result;
247
248    fn split_at(self, index: usize) -> (Self, Self, Self::Reducer) {
249        let (left, right, reducer) = self.base.split_at(index);
250        (
251            MapWithConsumer::new(left, self.item.clone(), self.map_op),
252            MapWithConsumer::new(right, self.item, self.map_op),
253            reducer,
254        )
255    }
256
257    fn into_folder(self) -> Self::Folder {
258        MapWithFolder {
259            base: self.base.into_folder(),
260            item: self.item,
261            map_op: self.map_op,
262        }
263    }
264
265    fn full(&self) -> bool {
266        self.base.full()
267    }
268}
269
270impl<'f, T, U, R, C, F> UnindexedConsumer<T> for MapWithConsumer<'f, C, U, F>
271where
272    C: UnindexedConsumer<R>,
273    U: Send + Clone,
274    F: Fn(&mut U, T) -> R + Sync,
275    R: Send,
276{
277    fn split_off_left(&self) -> Self {
278        MapWithConsumer::new(self.base.split_off_left(), self.item.clone(), self.map_op)
279    }
280
281    fn to_reducer(&self) -> Self::Reducer {
282        self.base.to_reducer()
283    }
284}
285
286struct MapWithFolder<'f, C, U, F> {
287    base: C,
288    item: U,
289    map_op: &'f F,
290}
291
292impl<'f, T, U, R, C, F> Folder<T> for MapWithFolder<'f, C, U, F>
293where
294    C: Folder<R>,
295    F: Fn(&mut U, T) -> R,
296{
297    type Result = C::Result;
298
299    fn consume(mut self, item: T) -> Self {
300        let mapped_item = (self.map_op)(&mut self.item, item);
301        self.base = self.base.consume(mapped_item);
302        self
303    }
304
305    fn consume_iter<I>(mut self, iter: I) -> Self
306    where
307        I: IntoIterator<Item = T>,
308    {
309        fn with<'f, T, U, R>(
310            item: &'f mut U,
311            map_op: impl Fn(&mut U, T) -> R + 'f,
312        ) -> impl FnMut(T) -> R + 'f {
313            move |x| map_op(item, x)
314        }
315
316        {
317            let mapped_iter = iter.into_iter().map(with(&mut self.item, self.map_op));
318            self.base = self.base.consume_iter(mapped_iter);
319        }
320        self
321    }
322
323    fn complete(self) -> C::Result {
324        self.base.complete()
325    }
326
327    fn full(&self) -> bool {
328        self.base.full()
329    }
330}
331
332// ------------------------------------------------------------------------------------------------
333
334/// `MapInit` is an iterator that transforms the elements of an underlying iterator.
335///
336/// This struct is created by the [`map_init()`] method on [`ParallelIterator`]
337///
338/// [`map_init()`]: ParallelIterator::map_init()
339#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
340#[derive(Clone)]
341pub struct MapInit<I, INIT, F> {
342    base: I,
343    init: INIT,
344    map_op: F,
345}
346
347impl<I: Debug, INIT, F> Debug for MapInit<I, INIT, F> {
348    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
349        f.debug_struct("MapInit").field("base", &self.base).finish()
350    }
351}
352
353impl<I, INIT, F> MapInit<I, INIT, F> {
354    /// Creates a new `MapInit` iterator.
355    pub(super) fn new(base: I, init: INIT, map_op: F) -> Self {
356        MapInit { base, init, map_op }
357    }
358}
359
360impl<I, INIT, T, F, R> ParallelIterator for MapInit<I, INIT, F>
361where
362    I: ParallelIterator,
363    INIT: Fn() -> T + Sync + Send,
364    F: Fn(&mut T, I::Item) -> R + Sync + Send,
365    R: Send,
366{
367    type Item = R;
368
369    fn drive_unindexed<C>(self, consumer: C) -> C::Result
370    where
371        C: UnindexedConsumer<Self::Item>,
372    {
373        let consumer1 = MapInitConsumer::new(consumer, &self.init, &self.map_op);
374        self.base.drive_unindexed(consumer1)
375    }
376
377    fn opt_len(&self) -> Option<usize> {
378        self.base.opt_len()
379    }
380}
381
382impl<I, INIT, T, F, R> IndexedParallelIterator for MapInit<I, INIT, F>
383where
384    I: IndexedParallelIterator,
385    INIT: Fn() -> T + Sync + Send,
386    F: Fn(&mut T, I::Item) -> R + Sync + Send,
387    R: Send,
388{
389    fn drive<C>(self, consumer: C) -> C::Result
390    where
391        C: Consumer<Self::Item>,
392    {
393        let consumer1 = MapInitConsumer::new(consumer, &self.init, &self.map_op);
394        self.base.drive(consumer1)
395    }
396
397    fn len(&self) -> usize {
398        self.base.len()
399    }
400
401    fn with_producer<CB>(self, callback: CB) -> CB::Output
402    where
403        CB: ProducerCallback<Self::Item>,
404    {
405        return self.base.with_producer(Callback {
406            callback,
407            init: self.init,
408            map_op: self.map_op,
409        });
410
411        struct Callback<CB, INIT, F> {
412            callback: CB,
413            init: INIT,
414            map_op: F,
415        }
416
417        impl<T, INIT, U, F, R, CB> ProducerCallback<T> for Callback<CB, INIT, F>
418        where
419            CB: ProducerCallback<R>,
420            INIT: Fn() -> U + Sync,
421            F: Fn(&mut U, T) -> R + Sync,
422            R: Send,
423        {
424            type Output = CB::Output;
425
426            fn callback<P>(self, base: P) -> CB::Output
427            where
428                P: Producer<Item = T>,
429            {
430                let producer = MapInitProducer {
431                    base,
432                    init: &self.init,
433                    map_op: &self.map_op,
434                };
435                self.callback.callback(producer)
436            }
437        }
438    }
439}
440
441// ////////////////////////////////////////////////////////////////////////
442
443struct MapInitProducer<'f, P, INIT, F> {
444    base: P,
445    init: &'f INIT,
446    map_op: &'f F,
447}
448
449impl<'f, P, INIT, U, F, R> Producer for MapInitProducer<'f, P, INIT, F>
450where
451    P: Producer,
452    INIT: Fn() -> U + Sync,
453    F: Fn(&mut U, P::Item) -> R + Sync,
454    R: Send,
455{
456    type Item = R;
457    type IntoIter = MapWithIter<'f, P::IntoIter, U, F>;
458
459    fn into_iter(self) -> Self::IntoIter {
460        MapWithIter {
461            base: self.base.into_iter(),
462            item: (self.init)(),
463            map_op: self.map_op,
464        }
465    }
466
467    fn min_len(&self) -> usize {
468        self.base.min_len()
469    }
470    fn max_len(&self) -> usize {
471        self.base.max_len()
472    }
473
474    fn split_at(self, index: usize) -> (Self, Self) {
475        let (left, right) = self.base.split_at(index);
476        (
477            MapInitProducer {
478                base: left,
479                init: self.init,
480                map_op: self.map_op,
481            },
482            MapInitProducer {
483                base: right,
484                init: self.init,
485                map_op: self.map_op,
486            },
487        )
488    }
489
490    fn fold_with<G>(self, folder: G) -> G
491    where
492        G: Folder<Self::Item>,
493    {
494        let folder1 = MapWithFolder {
495            base: folder,
496            item: (self.init)(),
497            map_op: self.map_op,
498        };
499        self.base.fold_with(folder1).base
500    }
501}
502
503// ////////////////////////////////////////////////////////////////////////
504// Consumer implementation
505
506struct MapInitConsumer<'f, C, INIT, F> {
507    base: C,
508    init: &'f INIT,
509    map_op: &'f F,
510}
511
512impl<'f, C, INIT, F> MapInitConsumer<'f, C, INIT, F> {
513    fn new(base: C, init: &'f INIT, map_op: &'f F) -> Self {
514        MapInitConsumer { base, init, map_op }
515    }
516}
517
518impl<'f, T, INIT, U, R, C, F> Consumer<T> for MapInitConsumer<'f, C, INIT, F>
519where
520    C: Consumer<R>,
521    INIT: Fn() -> U + Sync,
522    F: Fn(&mut U, T) -> R + Sync,
523    R: Send,
524{
525    type Folder = MapWithFolder<'f, C::Folder, U, F>;
526    type Reducer = C::Reducer;
527    type Result = C::Result;
528
529    fn split_at(self, index: usize) -> (Self, Self, Self::Reducer) {
530        let (left, right, reducer) = self.base.split_at(index);
531        (
532            MapInitConsumer::new(left, self.init, self.map_op),
533            MapInitConsumer::new(right, self.init, self.map_op),
534            reducer,
535        )
536    }
537
538    fn into_folder(self) -> Self::Folder {
539        MapWithFolder {
540            base: self.base.into_folder(),
541            item: (self.init)(),
542            map_op: self.map_op,
543        }
544    }
545
546    fn full(&self) -> bool {
547        self.base.full()
548    }
549}
550
551impl<'f, T, INIT, U, R, C, F> UnindexedConsumer<T> for MapInitConsumer<'f, C, INIT, F>
552where
553    C: UnindexedConsumer<R>,
554    INIT: Fn() -> U + Sync,
555    F: Fn(&mut U, T) -> R + Sync,
556    R: Send,
557{
558    fn split_off_left(&self) -> Self {
559        MapInitConsumer::new(self.base.split_off_left(), self.init, self.map_op)
560    }
561
562    fn to_reducer(&self) -> Self::Reducer {
563        self.base.to_reducer()
564    }
565}