compare_dir/
progress.rs

1use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
2use indicatif_log_bridge::LogWrapper;
3use std::io::{IsTerminal, stderr};
4use std::path::Path;
5use std::sync::atomic::{self, AtomicBool};
6use std::sync::{Arc, Mutex, Weak};
7use std::time::Duration;
8
9const SPINNER_STYLE0: &str = "{elapsed_precise} {spinner:.green} ";
10const SPINNER_STYLE1_NUM: &str = "{pos:>7} {msg}";
11const SPINNER_STYLE1_SIZE: &str = "{bytes:>7} {msg}";
12const NORMAL_STYLE0: &str = "{elapsed_precise} +{eta:>3} {percent:>3}% {bar:40.cyan/blue} ";
13const NORMAL_STYLE1_NUM: &str = "{pos:>7}/{len:7} {msg}";
14const NORMAL_STYLE1_SIZE: &str = "{bytes:>7}/{total_bytes:7} {msg}";
15const FILE_STYLE: &str = "  {elapsed:>3} +{eta:>3} {percent:>3}% {bar:10.cyan/blue} {wide_msg}";
16
17#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
18pub(crate) struct ProgressValue {
19    pub(crate) num_files: u64,
20    size: u64,
21}
22
23impl ProgressValue {
24    pub(crate) fn with_file_and_size(num_files: u64, size: u64) -> Self {
25        Self { size, num_files }
26    }
27
28    pub(crate) fn with_size(size: u64) -> Self {
29        Self::with_file_and_size(1, size)
30    }
31
32    fn is_zero(&self) -> bool {
33        self.num_files == 0 && self.size == 0
34    }
35
36    fn saturating_sub(&self, other: &ProgressValue) -> Self {
37        Self {
38            num_files: self.num_files.saturating_sub(other.num_files),
39            size: self.size.saturating_sub(other.size),
40        }
41    }
42
43    fn assert_valid_for_len(&self, len: &ProgressValue) {
44        assert!(
45            self.num_files <= len.num_files && self.size <= len.size,
46            "pos: {self:?}, len: {len:?}"
47        );
48    }
49}
50
51impl std::ops::Add for ProgressValue {
52    type Output = Self;
53
54    fn add(self, other: Self) -> Self {
55        Self {
56            num_files: self.num_files + other.num_files,
57            size: self.size + other.size,
58        }
59    }
60}
61
62impl std::ops::AddAssign for ProgressValue {
63    fn add_assign(&mut self, other: Self) {
64        self.num_files += other.num_files;
65        self.size += other.size;
66    }
67}
68
69#[derive(Debug, Default)]
70pub(crate) struct Progress {
71    inner: Option<ProgressBar>,
72    pos: ProgressValue,
73    use_bytes: bool,
74    is_finished: bool,
75    len: Option<ProgressValue>,
76    multi: Option<MultiProgress>,
77    primary: Weak<Mutex<Progress>>,
78}
79
80impl Drop for Progress {
81    fn drop(&mut self) {
82        if !self.is_finished {
83            self._inc_remaining();
84            self.finish();
85        }
86    }
87}
88
89impl Progress {
90    pub fn none() -> Self {
91        Self::none_with_primary(Weak::new())
92    }
93
94    fn none_with_primary(primary: Weak<Mutex<Progress>>) -> Self {
95        Self {
96            inner: None,
97            multi: None,
98            primary,
99            ..Default::default()
100        }
101    }
102
103    fn update_style(&self) {
104        if let Some(inner) = &self.inner {
105            let style = if self.len.is_some() {
106                match self.use_bytes {
107                    true => format!("{NORMAL_STYLE0}{NORMAL_STYLE1_SIZE}"),
108                    false => format!("{NORMAL_STYLE0}{NORMAL_STYLE1_NUM}"),
109                }
110            } else {
111                match self.use_bytes {
112                    true => format!("{SPINNER_STYLE0}{SPINNER_STYLE1_SIZE}"),
113                    false => format!("{SPINNER_STYLE0}{SPINNER_STYLE1_NUM}"),
114                }
115            };
116            inner.set_style(ProgressStyle::with_template(&style).unwrap());
117        }
118    }
119
120    fn update_position(&self) {
121        if let Some(inner) = &self.inner {
122            inner.set_position(match self.use_bytes {
123                true => self.pos.size,
124                false => self.pos.num_files,
125            });
126        }
127    }
128
129    fn update_length(&self) {
130        if let Some(inner) = &self.inner
131            && let Some(len) = self.len
132        {
133            inner.set_length(match self.use_bytes {
134                true => len.size,
135                false => len.num_files,
136            });
137        }
138    }
139
140    pub fn set_message(&self, msg: impl Into<String>) {
141        if let Some(inner) = &self.inner {
142            inner.set_message(msg.into());
143        }
144    }
145
146    pub fn inc(&mut self, amount: ProgressValue) {
147        self.pos += amount;
148        self.assert_pos();
149        self.update_position();
150        if let Some(primary) = self.primary.upgrade() {
151            primary.lock().unwrap().inc(amount);
152        }
153    }
154
155    pub fn inc_size(&mut self, size: u64) {
156        let amount = ProgressValue::with_file_and_size(0, size);
157        self.inc(amount);
158    }
159
160    pub fn inc_file(&mut self, num_files: u64) {
161        let amount = ProgressValue::with_file_and_size(num_files, 0);
162        self.inc(amount);
163    }
164
165    fn _inc_remaining(&mut self) {
166        self.assert_pos();
167        if let Some(len) = self.len {
168            let remaining = len.saturating_sub(&self.pos);
169            if !remaining.is_zero() {
170                log::debug!("remaining: {remaining:?}");
171                self.pos = len;
172                if let Some(primary) = self.primary.upgrade() {
173                    primary.lock().unwrap().inc(remaining);
174                }
175            }
176        }
177    }
178
179    pub fn set_length(&mut self, len: ProgressValue) {
180        self.len = Some(len);
181        self.update_style();
182        self.update_length();
183    }
184
185    fn assert_pos(&self) {
186        if let Some(len) = self.len {
187            self.pos.assert_valid_for_len(&len);
188        }
189    }
190
191    pub fn use_bytes(&mut self) {
192        self.use_bytes = true;
193        self.update_style();
194        self.update_length();
195        self.update_position();
196    }
197
198    pub fn finish(&mut self) {
199        assert!(!self.is_finished);
200        self.is_finished = true;
201        if let Some(len) = self.len {
202            assert_eq!(self.pos, len);
203        }
204        if let Some(inner) = &self.inner {
205            inner.finish();
206            if let Some(multi) = &self.multi {
207                multi.remove(inner);
208            }
209        }
210    }
211
212    pub fn suspend_for<F, R, S: IsTerminal>(&self, stream: S, f: F) -> R
213    where
214        F: FnOnce() -> R,
215    {
216        if !stream.is_terminal() {
217            return f();
218        }
219        if let Some(multi) = &self.multi {
220            multi.suspend(f)
221        } else if let Some(inner) = &self.inner {
222            inner.suspend(f)
223        } else {
224            f()
225        }
226    }
227}
228
229#[derive(Clone)]
230pub(crate) struct SharedProgress {
231    inner: Arc<Mutex<Progress>>,
232}
233
234impl SharedProgress {
235    pub fn none() -> Self {
236        Self {
237            inner: Arc::new(Mutex::new(Progress::none())),
238        }
239    }
240
241    pub fn set_message(&self, msg: impl Into<String>) {
242        self.inner.lock().unwrap().set_message(msg);
243    }
244
245    pub fn inc(&self, amount: ProgressValue) {
246        self.inner.lock().unwrap().inc(amount);
247    }
248
249    pub fn set_length(&self, len: ProgressValue) {
250        self.inner.lock().unwrap().set_length(len);
251    }
252
253    pub fn use_bytes(&self) {
254        self.inner.lock().unwrap().use_bytes();
255    }
256
257    pub fn finish(&self) {
258        self.inner.lock().unwrap().finish();
259    }
260
261    pub fn suspend_for<F, R, S: IsTerminal>(&self, stream: S, f: F) -> R
262    where
263        F: FnOnce() -> R,
264    {
265        self.inner.lock().unwrap().suspend_for(stream, f)
266    }
267}
268
269#[derive(Debug)]
270pub struct ProgressBuilder {
271    multi: MultiProgress,
272    pub is_enabled: bool,
273    pub is_file_enabled: bool,
274    is_propagate: AtomicBool,
275    primary: Mutex<Weak<Mutex<Progress>>>,
276}
277
278impl Default for ProgressBuilder {
279    fn default() -> Self {
280        Self {
281            multi: MultiProgress::default(),
282            is_enabled: stderr().is_terminal(),
283            is_file_enabled: false,
284            is_propagate: AtomicBool::new(false),
285            primary: Mutex::new(Weak::new()),
286        }
287    }
288}
289
290impl ProgressBuilder {
291    pub fn new() -> Self {
292        Self::default()
293    }
294
295    pub fn init_logger(&self, logger: env_logger::Logger) -> anyhow::Result<()> {
296        let max_level = logger.filter();
297        LogWrapper::new(self.multi.clone(), logger).try_init()?;
298        log::set_max_level(max_level);
299        Ok(())
300    }
301
302    pub(crate) fn is_propagate(&self) -> bool {
303        self.is_propagate.load(atomic::Ordering::Relaxed)
304    }
305
306    pub(crate) fn set_propagate(&self) {
307        self.is_propagate.store(true, atomic::Ordering::Relaxed);
308    }
309
310    pub(crate) fn add_primary(&self) -> SharedProgress {
311        if !self.is_enabled {
312            return SharedProgress::none();
313        }
314        let inner = self.multi.add(ProgressBar::new_spinner());
315        inner.enable_steady_tick(Duration::from_secs(1));
316        let progress = Progress {
317            inner: Some(inner),
318            multi: Some(self.multi.clone()),
319            primary: Weak::new(),
320            ..Default::default()
321        };
322        progress.update_style();
323        let shared = SharedProgress {
324            inner: Arc::new(Mutex::new(progress)),
325        };
326        *self.primary.lock().unwrap() = Arc::downgrade(&shared.inner);
327        shared
328    }
329
330    pub(crate) fn add_file(&self, path: &Path, file_size: u64) -> Progress {
331        if !self.is_enabled {
332            return Progress::none();
333        }
334        let primary = if self.is_propagate() {
335            Weak::clone(&self.primary.lock().unwrap())
336        } else {
337            Weak::new()
338        };
339        if !self.is_file_enabled {
340            return Progress::none_with_primary(primary);
341        }
342
343        let inner = self.multi.add(ProgressBar::new(file_size));
344        inner.set_style(ProgressStyle::with_template(FILE_STYLE).unwrap());
345        if let Some(parent) = path.parent()
346            && let Some(file_name) = path.file_name()
347        {
348            inner.set_message(format!(
349                "{} ({})",
350                file_name.to_string_lossy(),
351                parent.to_string_lossy()
352            ));
353        } else {
354            inner.set_message(path.to_string_lossy().to_string());
355        }
356        Progress {
357            inner: Some(inner),
358            multi: Some(self.multi.clone()),
359            use_bytes: true,
360            len: Some(ProgressValue::with_size(file_size)),
361            primary,
362            ..Default::default()
363        }
364    }
365}
366
367#[cfg(test)]
368mod tests {
369    use super::*;
370
371    #[test]
372    fn propagate() {
373        let builder = ProgressBuilder {
374            is_enabled: true,
375            is_file_enabled: false,
376            ..Default::default()
377        };
378        builder.set_propagate();
379        let _primary = builder.add_primary();
380
381        // Add a file progress
382        let file_path = Path::new("dummy.txt");
383        let mut file_progress = builder.add_file(file_path, 100);
384
385        // Initially primary position is 0
386        {
387            let prim = builder.primary.lock().unwrap().upgrade().unwrap();
388            assert_eq!(prim.lock().unwrap().pos.size, 0);
389        }
390
391        // Increment file progress
392        file_progress.inc(ProgressValue::with_size(10));
393
394        // Primary progress should be incremented to 10
395        {
396            let prim = builder.primary.lock().unwrap().upgrade().unwrap();
397            assert_eq!(prim.lock().unwrap().pos.size, 10);
398            assert_eq!(prim.lock().unwrap().pos.num_files, 1);
399        }
400    }
401}
compare_dir/progress.rs

compare_dir/
progress.rs
