1 use crate::progress::ProgressBar;
2 
3 /// Wraps an iterator to display its progress.
4 pub trait ProgressIterator
5 where
6     Self: Sized + Iterator,
7 {
8     /// Wrap an iterator with default styling. Attempt to guess iterator
9     /// length using `Iterator::size_hint`.
progress(self) -> ProgressBarIter<Self>10     fn progress(self) -> ProgressBarIter<Self> {
11         let n = match self.size_hint() {
12             (_, Some(n)) => n as u64,
13             _ => 0,
14         };
15         self.progress_count(n)
16     }
17 
18     /// Wrap an iterator with an explicit element count.
progress_count(self, len: u64) -> ProgressBarIter<Self>19     fn progress_count(self, len: u64) -> ProgressBarIter<Self> {
20         self.progress_with(ProgressBar::new(len))
21     }
22 
23     /// Wrap an iterator with a custom progress bar.
progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self>24     fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self>;
25 }
26 
27 /// Wraps an iterator to display its progress.
28 pub struct ProgressBarIter<T> {
29     it: T,
30     progress: ProgressBar,
31 }
32 
33 impl<S, T: Iterator<Item = S>> Iterator for ProgressBarIter<T> {
34     type Item = S;
35 
next(&mut self) -> Option<Self::Item>36     fn next(&mut self) -> Option<Self::Item> {
37         let next = self.it.next();
38 
39         if next.is_some() {
40             self.progress.inc(1);
41         } else {
42             self.progress.finish();
43         }
44 
45         next
46     }
47 }
48 
49 impl<S, T: Iterator<Item = S>> ProgressIterator for T {
progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self>50     fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self> {
51         ProgressBarIter { it: self, progress }
52     }
53 }
54 
55 #[cfg(feature = "with_rayon")]
56 pub mod rayon_support {
57     use super::*;
58     use rayon::iter::{
59         plumbing::Consumer, plumbing::Folder, plumbing::UnindexedConsumer, ParallelIterator,
60     };
61     use std::sync::{Arc, Mutex};
62 
63     pub struct ParProgressBarIter<T> {
64         it: T,
65         progress: Arc<Mutex<ProgressBar>>,
66     }
67 
68     /// Wraps a Rayon parallel iterator.
69     ///
70     /// See [`ProgressIterator`](trait.ProgressIterator.html) for method
71     /// documentation.
72     pub trait ParallelProgressIterator
73     where
74         Self: Sized,
75     {
progress_with(self, progress: ProgressBar) -> ParProgressBarIter<Self>76         fn progress_with(self, progress: ProgressBar) -> ParProgressBarIter<Self>;
77 
progress_count(self, len: u64) -> ParProgressBarIter<Self>78         fn progress_count(self, len: u64) -> ParProgressBarIter<Self> {
79             self.progress_with(ProgressBar::new(len))
80         }
81 
progress(self) -> ParProgressBarIter<Self>82         fn progress(self) -> ParProgressBarIter<Self> {
83             self.progress_count(0)
84         }
85     }
86 
87     impl<S: Send, T: ParallelIterator<Item = S>> ParallelProgressIterator for T {
progress_with(self, progress: ProgressBar) -> ParProgressBarIter<Self>88         fn progress_with(self, progress: ProgressBar) -> ParProgressBarIter<Self> {
89             ParProgressBarIter {
90                 it: self,
91                 progress: Arc::new(Mutex::new(progress)),
92             }
93         }
94     }
95 
96     struct ProgressConsumer<C> {
97         base: C,
98         progress: Arc<Mutex<ProgressBar>>,
99     }
100 
101     impl<C> ProgressConsumer<C> {
new(base: C, progress: Arc<Mutex<ProgressBar>>) -> Self102         fn new(base: C, progress: Arc<Mutex<ProgressBar>>) -> Self {
103             ProgressConsumer { base, progress }
104         }
105     }
106 
107     impl<T, C: Consumer<T>> Consumer<T> for ProgressConsumer<C> {
108         type Folder = ProgressFolder<C::Folder>;
109         type Reducer = C::Reducer;
110         type Result = C::Result;
111 
split_at(self, index: usize) -> (Self, Self, Self::Reducer)112         fn split_at(self, index: usize) -> (Self, Self, Self::Reducer) {
113             let (left, right, reducer) = self.base.split_at(index);
114             (
115                 ProgressConsumer::new(left, self.progress.clone()),
116                 ProgressConsumer::new(right, self.progress.clone()),
117                 reducer,
118             )
119         }
120 
into_folder(self) -> Self::Folder121         fn into_folder(self) -> Self::Folder {
122             ProgressFolder {
123                 base: self.base.into_folder(),
124                 progress: self.progress.clone(),
125             }
126         }
127 
full(&self) -> bool128         fn full(&self) -> bool {
129             self.base.full()
130         }
131     }
132 
133     impl<T, C: UnindexedConsumer<T>> UnindexedConsumer<T> for ProgressConsumer<C> {
split_off_left(&self) -> Self134         fn split_off_left(&self) -> Self {
135             ProgressConsumer::new(self.base.split_off_left(), self.progress.clone())
136         }
137 
to_reducer(&self) -> Self::Reducer138         fn to_reducer(&self) -> Self::Reducer {
139             self.base.to_reducer()
140         }
141     }
142 
143     struct ProgressFolder<C> {
144         base: C,
145         progress: Arc<Mutex<ProgressBar>>,
146     }
147 
148     impl<T, C: Folder<T>> Folder<T> for ProgressFolder<C> {
149         type Result = C::Result;
150 
consume(self, item: T) -> Self151         fn consume(self, item: T) -> Self {
152             self.progress.lock().unwrap().inc(1);
153             ProgressFolder {
154                 base: self.base.consume(item),
155                 progress: self.progress,
156             }
157         }
158 
complete(self) -> C::Result159         fn complete(self) -> C::Result {
160             self.base.complete()
161         }
162 
full(&self) -> bool163         fn full(&self) -> bool {
164             self.base.full()
165         }
166     }
167 
168     impl<S: Send, T: ParallelIterator<Item = S>> ParallelIterator for ParProgressBarIter<T> {
169         type Item = S;
170 
drive_unindexed<C: UnindexedConsumer<Self::Item>>(self, consumer: C) -> C::Result171         fn drive_unindexed<C: UnindexedConsumer<Self::Item>>(self, consumer: C) -> C::Result {
172             let consumer1 = ProgressConsumer::new(consumer, self.progress.clone());
173             self.it.drive_unindexed(consumer1)
174         }
175     }
176 
177     #[cfg(test)]
178     mod test {
179         use super::ParProgressBarIter;
180         use crate::iter::rayon_support::ParallelProgressIterator;
181         use crate::progress::ProgressBar;
182         use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
183 
184         #[test]
it_can_wrap_a_parallel_iterator()185         fn it_can_wrap_a_parallel_iterator() {
186             let v = vec![1, 2, 3];
187             let wrap = |it: ParProgressBarIter<_>| {
188                 assert_eq!(it.map(|x| x * 2).collect::<Vec<_>>(), vec![2, 4, 6]);
189             };
190 
191             wrap(v.par_iter().progress());
192             wrap(v.par_iter().progress_count(3));
193             wrap({
194                 let pb = ProgressBar::new(v.len() as u64);
195                 v.par_iter().progress_with(pb)
196             });
197         }
198     }
199 }
200 
201 #[cfg(test)]
202 mod test {
203     use crate::iter::{ProgressBarIter, ProgressIterator};
204     use crate::progress::ProgressBar;
205 
206     #[test]
it_can_wrap_an_iterator()207     fn it_can_wrap_an_iterator() {
208         let v = vec![1, 2, 3];
209         let wrap = |it: ProgressBarIter<_>| {
210             assert_eq!(it.map(|x| x * 2).collect::<Vec<_>>(), vec![2, 4, 6]);
211         };
212 
213         wrap(v.iter().progress());
214         wrap(v.iter().progress_count(3));
215         wrap({
216             let pb = ProgressBar::new(v.len() as u64);
217             v.iter().progress_with(pb)
218         });
219     }
220 }
221