1use super::plumbing::*;
2use super::*;
3use 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)]
12pub struct Take<I> {
13 base: I,
14 n: usize,
15}
16
17impl<I> Take<I>
18where
19 I: IndexedParallelIterator,
20{
21 /// Creates a new `Take` iterator.
22 pub(super) fn new(base: I, n: usize) -> Self {
23 let n: usize = min(v1:base.len(), v2:n);
24 Take { base, n }
25 }
26}
27
28impl<I> ParallelIterator for Take<I>
29where
30 I: IndexedParallelIterator,
31{
32 type Item = I::Item;
33
34 fn drive_unindexed<C>(self, consumer: C) -> C::Result
35 where
36 C: UnindexedConsumer<Self::Item>,
37 {
38 bridge(self, consumer)
39 }
40
41 fn opt_len(&self) -> Option<usize> {
42 Some(self.len())
43 }
44}
45
46impl<I> IndexedParallelIterator for Take<I>
47where
48 I: IndexedParallelIterator,
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
59 where
60 CB: ProducerCallback<Self::Item>,
61 {
62 return self.base.with_producer(Callback {
63 callback,
64 n: self.n,
65 });
66
67 struct Callback<CB> {
68 callback: CB,
69 n: usize,
70 }
71
72 impl<T, CB> ProducerCallback<T> for Callback<CB>
73 where
74 CB: ProducerCallback<T>,
75 {
76 type Output = CB::Output;
77 fn callback<P>(self, base: P) -> CB::Output
78 where
79 P: Producer<Item = T>,
80 {
81 let (producer, _) = base.split_at(self.n);
82 self.callback.callback(producer)
83 }
84 }
85 }
86}
87