1use std::{io::{stderr, Write}, fmt::Display, time::Instant, sync::atomic::{AtomicU64, Ordering::SeqCst}};
2
3#[cfg(feature = "num-format")]
4use num_format::{Locale, ToFormattedString, ToFormattedStr};
5
6#[derive(Clone)]
7pub enum Style {
8 Mono(char),
9 Edged(char, char),
10}
11
12impl Style {
13 fn bar_char(&self) -> char {
14 match *self { Self::Mono(c) | Self::Edged(c, _) => c }
15 }
16
17 fn edge_char(&self) -> char {
18 match *self { Self::Mono(c) | Self::Edged(_, c) => c }
19 }
20}
21
22#[derive(Clone)]
23pub struct Config<'a> {
24 pub width: Option<u64>,
25 pub default_width: u64,
26 pub delimiters: (char, char),
27 pub style: Style,
28 pub space_char: char,
29 pub prefix: &'a str,
30 pub unit: &'a str,
31 pub num_width: usize,
32 pub throttle_millis: u64,
33}
34
35impl Config<'_> {
36 #[inline]
37 pub fn ascii() -> Self {
38 Self { style: Style::Mono('#'), ..Default::default() }
39 }
40
41 #[inline]
42 pub fn unicode() -> Self {
43 Self { style: Style::Mono('█'), ..Default::default() }
44 }
45
46 #[inline]
47 pub fn cargo() -> Self {
48 Self { style: Style::Edged('=', '>'), ..Default::default() }
49 }
50}
51
52impl Default for Config<'_> {
53 fn default() -> Self {
54 Self {
55 width: None,
56 default_width: 80,
57 delimiters: ('[', ']'),
58 style: Style::Mono('#'),
59 space_char: ' ',
60 prefix: "",
61 unit: "",
62 num_width: 0,
63 throttle_millis: 10,
64 }
65 }
66}
67
68#[inline]
69pub fn bar<I: ExactSizeIterator>(iter: I) -> impl Iterator<Item = I::Item> {
70 bar_with_config(iter, Config::default())
71}
72
73#[inline]
74pub fn bar_with_config<I: ExactSizeIterator>(iter: I, config: Config) -> std::iter::Inspect<I, impl FnMut(&I::Item) + '_> {
75 let bar = Bar::new(iter.len().try_into().unwrap(), config);
76 iter.inspect(move |_| bar.inc(1))
77}
78
79#[inline]
80pub fn bar_chunks<T>(chunk_size: usize, slice: &[T]) -> impl Iterator<Item = &T> {
81 bar_chunks_with_config(chunk_size, slice, Config::default())
82}
83
84#[inline]
85pub fn bar_chunks_with_config<'b, 'a: 'b, T>(chunk_size: usize, slice: &'a [T], config: Config<'b>) -> impl Iterator<Item = &'a T> + 'b {
86 let bar = Bar::new(slice.len().try_into().unwrap(), config);
87 slice.chunks(chunk_size).inspect(move |chunk| bar.inc(chunk.len() as u64)).flatten()
88}
89
90#[inline]
91pub fn bar_chunks_mut<T>(chunk_size: usize, slice: &mut [T]) -> impl Iterator<Item = &mut T> {
92 bar_chunks_mut_with_config(chunk_size, slice, Config::default())
93}
94
95#[inline]
96pub fn bar_chunks_mut_with_config<'b, 'a: 'b, T>(chunk_size: usize, slice: &'a mut [T], config: Config<'b>) -> impl Iterator<Item = &'a mut T> + 'b {
97 let bar = Bar::new(slice.len().try_into().unwrap(), config);
98 slice.chunks_mut(chunk_size).inspect(move |chunk| bar.inc(chunk.len() as u64)).flatten()
99}
100
101pub struct Bar<'a> {
102 config: Config<'a>,
103 len: u64,
104 pos: AtomicU64,
105 len_str: String,
106 bar_width: u64,
107 start_time: Instant,
108 last_update: AtomicU64,
109}
110
111impl<'a> Bar<'a> {
112 #[inline]
113 pub fn new(len: u64, mut config: Config<'a>) -> Self {
114 let len_str = format_number(len);
115 config.num_width = config.num_width.max(len_str.len());
116 #[cfg(feature = "terminal_size")]
117 { config.width = config.width.or_else(|| Some(u64::from(terminal_size::terminal_size()?.0.0))) }
118 let bar_width = config.width.unwrap_or(config.default_width) - 35 - (config.prefix.len() + config.unit.len() + config.num_width * 2) as u64
119 - if config.unit.is_empty() { 0 } else { 1 };
120 Self { config, bar_width, len, pos: AtomicU64::new(0), len_str, start_time: Instant::now(), last_update: AtomicU64::new(0) }
121 }
122
123 fn print(&self) -> std::io::Result<()> {
124 let mut stderr = stderr().lock();
125 let pos = self.pos.load(SeqCst);
126 assert!(pos <= self.len);
127 let ratio = (pos as f64) / (self.len as f64);
128 let progress_width = (ratio * (self.bar_width as f64)).round() as u64;
129 let secs_per_step = self.start_time.elapsed().as_secs_f64() / (pos as f64);
130 let eta = Time(((self.len.saturating_sub(pos) as f64) * secs_per_step).ceil() as u64);
131
132 write!(stderr, "\r{} {} {:>num_width$} / {:>num_width$}{}{} {}", self.config.prefix, Time(self.start_time.elapsed().as_secs()), format_number(pos),
133 self.len_str, if self.config.unit.is_empty() { "" } else { " " }, self.config.unit, self.config.delimiters.0, num_width = self.config.num_width)?;
134 write_iter(&mut stderr, std::iter::repeat(self.config.style.bar_char()).take(progress_width as usize))?;
135 write!(stderr, "{}", if pos == self.len { self.config.style.bar_char() } else { self.config.style.edge_char() })?;
136 write_iter(&mut stderr, std::iter::repeat(self.config.space_char).take((self.bar_width - progress_width) as usize))?;
137 write!(stderr, "{} {:3.0}% ETA {eta}\r", self.config.delimiters.1, ratio * 100.)?;
138 stderr.flush()?;
139 Ok(())
140 }
141
142 #[inline]
143 pub fn inc(&self, delta: u64) {
144 self.pos.fetch_add(delta, SeqCst);
145 let elapsed = self.elapsed_millis();
146 let last_update = self.last_update.load(SeqCst);
147
148 if elapsed - last_update > self.config.throttle_millis && self.last_update.compare_exchange(last_update, elapsed, SeqCst, SeqCst).is_ok() {
149 self.print().unwrap();
150 }
151 }
152
153 #[inline]
154 pub fn finish(self) {
155 drop(self);
156 }
157
158 fn elapsed_millis(&self) -> u64 {
159 self.start_time.elapsed().as_millis().try_into().unwrap()
160 }
161}
162
163impl Drop for Bar<'_> {
164 #[inline]
165 fn drop(&mut self) {
166 self.print().unwrap();
167 eprintln!();
168 }
169}
170
171fn write_iter<W, I>(w: &mut W, mut iter: I) -> std::io::Result<()>
172where
173 W: Write,
174 I: Iterator,
175 I::Item: Display,
176{
177 iter.try_for_each(|x| write!(w, "{x}"))
178}
179
180#[cfg(feature = "num-format")]
181fn format_number<T: ToFormattedStr>(number: T) -> String {
182 number.to_formatted_string(&Locale::en)
183}
184
185#[cfg(not(feature = "num-format"))]
186fn format_number<T: Display>(number: T) -> String {
187 number.to_string()
188}
189
190struct Time(u64);
191
192impl Display for Time {
193 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
194 let hours = self.0 / 3600;
195
196 if hours > 99 {
197 write!(f, "??:??:??")
198 } else {
199 let mins = (self.0 / 60) % 60;
200 let secs = self.0 % 60;
201 write!(f, "{hours:02}:{mins:02}:{secs:02}")
202 }
203 }
204}