]>
Commit | Line | Data |
---|---|---|
2c00a5a8 XL |
1 | use super::plumbing::*; |
2 | use super::*; | |
3 | use std::cmp::min; | |
4 | ||
5 | /// `Take` is an iterator that iterates over the first `n` elements. | |
6 | /// This struct is created by the [`take()`] method on [`IndexedParallelIterator`] | |
7 | /// | |
8 | /// [`take()`]: trait.IndexedParallelIterator.html#method.take | |
9 | /// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html | |
10 | #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] | |
11 | #[derive(Debug, Clone)] | |
12 | pub struct Take<I> { | |
13 | base: I, | |
14 | n: usize, | |
15 | } | |
16 | ||
e74abb32 | 17 | impl<I> Take<I> |
532ac7d7 XL |
18 | where |
19 | I: IndexedParallelIterator, | |
2c00a5a8 | 20 | { |
923072b8 | 21 | /// Creates a new `Take` iterator. |
e74abb32 XL |
22 | pub(super) fn new(base: I, n: usize) -> Self { |
23 | let n = min(base.len(), n); | |
24 | Take { base, n } | |
25 | } | |
2c00a5a8 XL |
26 | } |
27 | ||
28 | impl<I> ParallelIterator for Take<I> | |
532ac7d7 XL |
29 | where |
30 | I: IndexedParallelIterator, | |
2c00a5a8 XL |
31 | { |
32 | type Item = I::Item; | |
33 | ||
34 | fn drive_unindexed<C>(self, consumer: C) -> C::Result | |
532ac7d7 XL |
35 | where |
36 | C: UnindexedConsumer<Self::Item>, | |
2c00a5a8 XL |
37 | { |
38 | bridge(self, consumer) | |
39 | } | |
40 | ||
41 | fn opt_len(&self) -> Option<usize> { | |
42 | Some(self.len()) | |
43 | } | |
44 | } | |
45 | ||
46 | impl<I> IndexedParallelIterator for Take<I> | |
532ac7d7 XL |
47 | where |
48 | I: IndexedParallelIterator, | |
2c00a5a8 XL |
49 | { |
50 | fn len(&self) -> usize { | |
51 | self.n | |
52 | } | |
53 | ||
54 | fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result { | |
55 | bridge(self, consumer) | |
56 | } | |
57 | ||
58 | fn with_producer<CB>(self, callback: CB) -> CB::Output | |
532ac7d7 XL |
59 | where |
60 | CB: ProducerCallback<Self::Item>, | |
2c00a5a8 XL |
61 | { |
62 | return self.base.with_producer(Callback { | |
e74abb32 | 63 | callback, |
532ac7d7 XL |
64 | n: self.n, |
65 | }); | |
2c00a5a8 XL |
66 | |
67 | struct Callback<CB> { | |
68 | callback: CB, | |
69 | n: usize, | |
70 | } | |
71 | ||
72 | impl<T, CB> ProducerCallback<T> for Callback<CB> | |
532ac7d7 XL |
73 | where |
74 | CB: ProducerCallback<T>, | |
2c00a5a8 XL |
75 | { |
76 | type Output = CB::Output; | |
77 | fn callback<P>(self, base: P) -> CB::Output | |
532ac7d7 XL |
78 | where |
79 | P: Producer<Item = T>, | |
2c00a5a8 XL |
80 | { |
81 | let (producer, _) = base.split_at(self.n); | |
82 | self.callback.callback(producer) | |
83 | } | |
84 | } | |
85 | } | |
86 | } |