indicatif/
iter.rs

1use std::borrow::Cow;
2use std::io::{self, IoSliceMut};
3use std::iter::FusedIterator;
4#[cfg(feature = "tokio")]
5use std::pin::Pin;
6#[cfg(feature = "tokio")]
7use std::task::{Context, Poll};
8use std::time::Duration;
9
10#[cfg(feature = "tokio")]
11use tokio::io::{ReadBuf, SeekFrom};
12
13use crate::progress_bar::ProgressBar;
14use crate::state::ProgressFinish;
15use crate::style::ProgressStyle;
16
17/// Wraps an iterator to display its progress.
18pub trait ProgressIterator
19where
20    Self: Sized + Iterator,
21{
22    /// Wrap an iterator with default styling. Uses [`Iterator::size_hint()`] to get length.
23    /// Returns `Some(..)` only if `size_hint.1` is [`Some`]. If you want to create a progress bar
24    /// even if `size_hint.1` returns [`None`] use [`progress_count()`](ProgressIterator::progress_count)
25    /// or [`progress_with()`](ProgressIterator::progress_with) instead.
26    fn try_progress(self) -> Option<ProgressBarIter<Self>> {
27        self.size_hint()
28            .1
29            .map(|len| self.progress_count(u64::try_from(len).unwrap()))
30    }
31
32    /// Wrap an iterator with default styling.
33    fn progress(self) -> ProgressBarIter<Self>
34    where
35        Self: ExactSizeIterator,
36    {
37        let len = u64::try_from(self.len()).unwrap();
38        self.progress_count(len)
39    }
40
41    /// Wrap an iterator with an explicit element count.
42    fn progress_count(self, len: u64) -> ProgressBarIter<Self> {
43        self.progress_with(ProgressBar::new(len))
44    }
45
46    /// Wrap an iterator with a custom progress bar.
47    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self>;
48
49    /// Wrap an iterator with a progress bar and style it.
50    fn progress_with_style(self, style: crate::ProgressStyle) -> ProgressBarIter<Self>
51    where
52        Self: ExactSizeIterator,
53    {
54        let len = u64::try_from(self.len()).unwrap();
55        let bar = ProgressBar::new(len).with_style(style);
56        self.progress_with(bar)
57    }
58}
59
60/// Wraps an iterator to display its progress.
61#[derive(Debug)]
62pub struct ProgressBarIter<T> {
63    pub(crate) it: T,
64    pub progress: ProgressBar,
65    pub(crate) seek_max: SeekMax,
66}
67
68impl<T> ProgressBarIter<T> {
69    /// Builder-like function for setting underlying progress bar's style.
70    ///
71    /// See [`ProgressBar::with_style()`].
72    pub fn with_style(mut self, style: ProgressStyle) -> Self {
73        self.progress = self.progress.with_style(style);
74        self
75    }
76
77    /// Builder-like function for setting underlying progress bar's prefix.
78    ///
79    /// See [`ProgressBar::with_prefix()`].
80    pub fn with_prefix(mut self, prefix: impl Into<Cow<'static, str>>) -> Self {
81        self.progress = self.progress.with_prefix(prefix);
82        self
83    }
84
85    /// Builder-like function for setting underlying progress bar's message.
86    ///
87    /// See [`ProgressBar::with_message()`].
88    pub fn with_message(mut self, message: impl Into<Cow<'static, str>>) -> Self {
89        self.progress = self.progress.with_message(message);
90        self
91    }
92
93    /// Builder-like function for setting underlying progress bar's position.
94    ///
95    /// See [`ProgressBar::with_position()`].
96    pub fn with_position(mut self, position: u64) -> Self {
97        self.progress = self.progress.with_position(position);
98        self
99    }
100
101    /// Builder-like function for setting underlying progress bar's elapsed time.
102    ///
103    /// See [`ProgressBar::with_elapsed()`].
104    pub fn with_elapsed(mut self, elapsed: Duration) -> Self {
105        self.progress = self.progress.with_elapsed(elapsed);
106        self
107    }
108
109    /// Builder-like function for setting underlying progress bar's finish behavior.
110    ///
111    /// See [`ProgressBar::with_finish()`].
112    pub fn with_finish(mut self, finish: ProgressFinish) -> Self {
113        self.progress = self.progress.with_finish(finish);
114        self
115    }
116}
117
118impl<S, T: Iterator<Item = S>> Iterator for ProgressBarIter<T> {
119    type Item = S;
120
121    fn next(&mut self) -> Option<Self::Item> {
122        let item = self.it.next();
123
124        if item.is_some() {
125            self.progress.inc(1);
126        } else if !self.progress.is_finished() {
127            self.progress.finish_using_style();
128        }
129
130        item
131    }
132}
133
134impl<T: ExactSizeIterator> ExactSizeIterator for ProgressBarIter<T> {
135    fn len(&self) -> usize {
136        self.it.len()
137    }
138}
139
140impl<T: DoubleEndedIterator> DoubleEndedIterator for ProgressBarIter<T> {
141    fn next_back(&mut self) -> Option<Self::Item> {
142        let item = self.it.next_back();
143
144        if item.is_some() {
145            self.progress.inc(1);
146        } else if !self.progress.is_finished() {
147            self.progress.finish_using_style();
148        }
149
150        item
151    }
152}
153
154impl<T: FusedIterator> FusedIterator for ProgressBarIter<T> {}
155
156impl<R: io::Read> io::Read for ProgressBarIter<R> {
157    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
158        let inc = self.it.read(buf)?;
159        self.progress.set_position(
160            self.seek_max
161                .update_seq(self.progress.position(), inc as u64),
162        );
163        Ok(inc)
164    }
165
166    fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
167        let inc = self.it.read_vectored(bufs)?;
168        self.progress.set_position(
169            self.seek_max
170                .update_seq(self.progress.position(), inc as u64),
171        );
172        Ok(inc)
173    }
174
175    fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
176        let inc = self.it.read_to_string(buf)?;
177        self.progress.set_position(
178            self.seek_max
179                .update_seq(self.progress.position(), inc as u64),
180        );
181        Ok(inc)
182    }
183
184    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
185        self.it.read_exact(buf)?;
186        self.progress.set_position(
187            self.seek_max
188                .update_seq(self.progress.position(), buf.len() as u64),
189        );
190        Ok(())
191    }
192}
193
194impl<R: io::BufRead> io::BufRead for ProgressBarIter<R> {
195    fn fill_buf(&mut self) -> io::Result<&[u8]> {
196        self.it.fill_buf()
197    }
198
199    fn consume(&mut self, amt: usize) {
200        self.it.consume(amt);
201        self.progress.set_position(
202            self.seek_max
203                .update_seq(self.progress.position(), amt.try_into().unwrap()),
204        );
205    }
206}
207
208impl<S: io::Seek> io::Seek for ProgressBarIter<S> {
209    fn seek(&mut self, f: io::SeekFrom) -> io::Result<u64> {
210        self.it.seek(f).inspect(|&pos| {
211            if f != io::SeekFrom::Current(0) {
212                // this kind of seek is used to find the current position, but does not alter it
213                // generally equivalent to stream_position()
214                self.progress.set_position(self.seek_max.update_seek(pos));
215            }
216        })
217    }
218    // Pass this through to preserve optimizations that the inner I/O object may use here
219    // Also avoid sending a set_position update when the position hasn't changed
220    fn stream_position(&mut self) -> io::Result<u64> {
221        self.it.stream_position()
222    }
223}
224
225/// Calculates a more stable visual position from jittery seeks to show to the user.
226///
227/// Holds the maximum position encountered out of the last HISTORY read/write positions.
228/// Drops history when only sequential operations are performed RESET times in a row.
229#[derive(Debug, Default)]
230pub(crate) struct SeekMax<const RESET: u8 = 5, const HISTORY: usize = 10> {
231    buf: Option<(Box<MaxRingBuf<HISTORY>>, u8)>,
232}
233
234impl<const RESET: u8, const HISTORY: usize> SeekMax<RESET, HISTORY> {
235    fn update_seq(&mut self, prev_pos: u64, delta: u64) -> u64 {
236        let new_pos = prev_pos + delta;
237        let Some((buf, seq)) = &mut self.buf else {
238            return new_pos;
239        };
240
241        *seq += 1;
242        if *seq >= RESET {
243            self.buf = None;
244            return new_pos;
245        }
246
247        buf.update(new_pos);
248        buf.max()
249    }
250
251    fn update_seek(&mut self, newpos: u64) -> u64 {
252        let (b, seq) = self
253            .buf
254            .get_or_insert_with(|| (Box::new(MaxRingBuf::<HISTORY>::default()), 0));
255        *seq = 0;
256        b.update(newpos);
257        b.max()
258    }
259}
260
261/// Ring buffer that remembers the maximum contained value.
262#[derive(Debug)]
263struct MaxRingBuf<const HISTORY: usize = 10> {
264    history: [u64; HISTORY],
265    head: u8,    // must be < HISTORY
266    max_pos: u8, // must be < HISTORY
267}
268
269impl<const HISTORY: usize> MaxRingBuf<HISTORY> {
270    /// Updates internal bookkeeping to remember the maximum value
271    ///
272    /// Updates that overwrite the position the maximum was stored in with a smaller number do a
273    /// seek of the buffer, searching for the new maximum. This only happens on average each
274    /// 1 / HISTORY and has a cost of HISTORY, therefore amortizing to O(1).
275    ///
276    /// In case there is some linear increase with jitter, as expected in this specific use-case,
277    /// as long as there is one bigger update each HISTORY updates the scan is never triggered at all.
278    ///
279    /// Worst case would be linearly decreasing values, which is still O(1).
280    fn update(&mut self, new: u64) {
281        let head = usize::from(self.head) % self.history.len();
282        let max_pos = usize::from(self.max_pos) % self.history.len();
283        let prev_max = self.history[max_pos];
284        self.history[head] = new;
285
286        if new > prev_max {
287            // This is now the new maximum
288            self.max_pos = self.head;
289        } else if self.max_pos == self.head && new < prev_max {
290            // This was the maximum and may not be anymore
291            // do a linear seek to find the new maximum
292            let (idx, _val) = self
293                .history
294                .iter()
295                .enumerate()
296                .max_by_key(|(_, v)| *v)
297                .expect("array has fixded size > 0");
298            // invariant_m: idx is from an enumeration of history
299            self.max_pos = idx as u8;
300        }
301
302        self.head = (self.head + 1) % (self.history.len() as u8);
303    }
304
305    fn max(&self) -> u64 {
306        // exploit invariant_m to eliminate bounds checks & panic code path
307        self.history[self.max_pos as usize % self.history.len()]
308    }
309}
310
311impl<const HISTORY: usize> Default for MaxRingBuf<HISTORY> {
312    fn default() -> Self {
313        assert!(HISTORY <= u8::MAX.into());
314        assert!(HISTORY > 0);
315        Self {
316            history: [0; HISTORY],
317            head: 0,
318            max_pos: 0,
319        }
320    }
321}
322
323#[cfg(feature = "tokio")]
324#[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
325impl<W: tokio::io::AsyncWrite + Unpin> tokio::io::AsyncWrite for ProgressBarIter<W> {
326    fn poll_write(
327        mut self: Pin<&mut Self>,
328        cx: &mut Context<'_>,
329        buf: &[u8],
330    ) -> Poll<io::Result<usize>> {
331        Pin::new(&mut self.it).poll_write(cx, buf).map(|poll| {
332            poll.inspect(|&inc| {
333                let pos = self.progress.position();
334                let new = self.seek_max.update_seq(pos, inc as u64);
335                self.progress.set_position(new);
336            })
337        })
338    }
339
340    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
341        Pin::new(&mut self.it).poll_flush(cx)
342    }
343
344    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
345        Pin::new(&mut self.it).poll_shutdown(cx)
346    }
347}
348
349#[cfg(feature = "tokio")]
350#[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
351impl<W: tokio::io::AsyncRead + Unpin> tokio::io::AsyncRead for ProgressBarIter<W> {
352    fn poll_read(
353        mut self: Pin<&mut Self>,
354        cx: &mut Context<'_>,
355        buf: &mut ReadBuf<'_>,
356    ) -> Poll<io::Result<()>> {
357        let prev_len = buf.filled().len() as u64;
358        let poll = Pin::new(&mut self.it).poll_read(cx, buf);
359        if let Poll::Ready(_e) = &poll {
360            let inc = buf.filled().len() as u64 - prev_len;
361            let pos = self.progress.position();
362            let new = self.seek_max.update_seq(pos, inc);
363            self.progress.set_position(new);
364        }
365        poll
366    }
367}
368
369#[cfg(feature = "tokio")]
370#[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
371impl<W: tokio::io::AsyncSeek + Unpin> tokio::io::AsyncSeek for ProgressBarIter<W> {
372    fn start_seek(mut self: Pin<&mut Self>, position: SeekFrom) -> io::Result<()> {
373        Pin::new(&mut self.it).start_seek(position)
374    }
375
376    fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
377        let poll = Pin::new(&mut self.it).poll_complete(cx);
378        if let Poll::Ready(Ok(pos)) = &poll {
379            let new = self.seek_max.update_seek(*pos);
380            self.progress.set_position(new);
381        }
382
383        poll
384    }
385}
386
387#[cfg(feature = "tokio")]
388#[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
389impl<W: tokio::io::AsyncBufRead + Unpin + tokio::io::AsyncRead> tokio::io::AsyncBufRead
390    for ProgressBarIter<W>
391{
392    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
393        let this = self.get_mut();
394        Pin::new(&mut this.it).poll_fill_buf(cx)
395    }
396
397    fn consume(mut self: Pin<&mut Self>, amt: usize) {
398        Pin::new(&mut self.it).consume(amt);
399        let pos = self.progress.position();
400        let new = self.seek_max.update_seq(pos, amt as u64);
401        self.progress.set_position(new);
402    }
403}
404
405#[cfg(feature = "futures")]
406#[cfg_attr(docsrs, doc(cfg(feature = "futures")))]
407impl<S: futures_core::Stream + Unpin> futures_core::Stream for ProgressBarIter<S> {
408    type Item = S::Item;
409
410    fn poll_next(
411        self: std::pin::Pin<&mut Self>,
412        cx: &mut std::task::Context<'_>,
413    ) -> std::task::Poll<Option<Self::Item>> {
414        let this = self.get_mut();
415        let item = std::pin::Pin::new(&mut this.it).poll_next(cx);
416        match &item {
417            std::task::Poll::Ready(Some(_)) => this.progress.inc(1),
418            std::task::Poll::Ready(None) => this.progress.finish_using_style(),
419            std::task::Poll::Pending => {}
420        }
421        item
422    }
423}
424
425impl<W: io::Write> io::Write for ProgressBarIter<W> {
426    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
427        self.it.write(buf).inspect(|&inc| {
428            self.progress.set_position(
429                self.seek_max
430                    .update_seq(self.progress.position(), inc as u64),
431            );
432        })
433    }
434
435    fn write_vectored(&mut self, bufs: &[io::IoSlice]) -> io::Result<usize> {
436        self.it.write_vectored(bufs).inspect(|&inc| {
437            self.progress.set_position(
438                self.seek_max
439                    .update_seq(self.progress.position(), inc as u64),
440            );
441        })
442    }
443
444    fn flush(&mut self) -> io::Result<()> {
445        self.it.flush()
446    }
447
448    // write_fmt can not be captured with reasonable effort.
449    // as it uses write_all internally by default that should not be a problem.
450    // fn write_fmt(&mut self, fmt: fmt::Arguments) -> io::Result<()>;
451}
452
453impl<S, T: Iterator<Item = S>> ProgressIterator for T {
454    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self> {
455        ProgressBarIter {
456            it: self,
457            progress,
458            seek_max: SeekMax::default(),
459        }
460    }
461}
462
463#[cfg(test)]
464mod test {
465    use crate::iter::{ProgressBarIter, ProgressIterator};
466    use crate::progress_bar::ProgressBar;
467    use crate::ProgressStyle;
468
469    #[test]
470    fn it_can_wrap_an_iterator() {
471        let v = [1, 2, 3];
472        let wrap = |it: ProgressBarIter<_>| {
473            assert_eq!(it.map(|x| x * 2).collect::<Vec<_>>(), vec![2, 4, 6]);
474        };
475
476        wrap(v.iter().progress());
477        wrap(v.iter().progress_count(3));
478        wrap({
479            let pb = ProgressBar::new(v.len() as u64);
480            v.iter().progress_with(pb)
481        });
482        wrap({
483            let style = ProgressStyle::default_bar()
484                .template("{wide_bar:.red} {percent}/100%")
485                .unwrap();
486            v.iter().progress_with_style(style)
487        });
488    }
489
490    #[test]
491    fn test_max_ring_buf() {
492        use crate::iter::MaxRingBuf;
493        let mut max = MaxRingBuf::<10>::default();
494        max.update(100);
495        assert_eq!(max.max(), 100);
496        for i in 0..10 {
497            max.update(99 - i);
498        }
499        assert_eq!(max.max(), 99);
500    }
501}
indicatif/iter.rs

indicatif/
iter.rs
