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