wb_cache/test/simulation/
progress.rs

1pub mod traits;
2
3use std::fmt::Debug;
4use std::fmt::Display;
5use std::fmt::Write;
6use std::time::Instant;
7
8use console::Style;
9use fieldx::fxstruct;
10use indicatif::style::ProgressTracker;
11use indicatif::MultiProgress;
12use indicatif::ProgressBar;
13use indicatif::ProgressState;
14use indicatif::ProgressStyle;
15pub use traits::*;
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
18pub enum MsgType {
19    Debug,
20    Info,
21    Warn,
22    Error,
23}
24
25impl Display for MsgType {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        match self {
28            MsgType::Debug => write!(f, "DEBUG"),
29            MsgType::Info => write!(f, "INFO"),
30            MsgType::Warn => write!(f, "WARN"),
31            MsgType::Error => write!(f, "ERROR"),
32        }
33    }
34}
35
36pub enum POrder<'a> {
37    Before(Option<&'a ProgressBar>),
38    After(Option<&'a ProgressBar>),
39}
40
41pub enum PStyle {
42    Main,
43    Minor,
44    Message,
45    Custom(ProgressStyle),
46}
47
48impl Debug for PStyle {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        match self {
51            PStyle::Main => write!(f, "Main"),
52            PStyle::Minor => write!(f, "Minor"),
53            PStyle::Message => write!(f, "Message"),
54            PStyle::Custom(_) => write!(f, "Custom"),
55        }
56    }
57}
58
59struct PerSecFmt;
60
61impl ProgressTracker for PerSecFmt {
62    fn clone_box(&self) -> Box<dyn ProgressTracker> {
63        Box::new(PerSecFmt)
64    }
65
66    fn tick(&mut self, _state: &ProgressState, _now: Instant) {}
67
68    fn reset(&mut self, _state: &ProgressState, _now: Instant) {}
69
70    fn write(&self, state: &ProgressState, w: &mut dyn Write) {
71        // Writes, for example, "12.34/s"
72        write!(w, "{:.2}/s", state.per_sec()).unwrap();
73    }
74}
75
76#[fxstruct(new(off), sync, fallible(off, error(SimError)), builder)]
77pub struct ProgressUI {
78    #[fieldx(get(copy), default(false))]
79    quiet: bool,
80
81    #[fieldx(lazy, get, builder(off))]
82    multi_progress: Option<MultiProgress>,
83
84    #[fieldx(lazy, get(copy))]
85    user_attended: bool,
86
87    #[fieldx(lazy, private, get(clone))]
88    progress_style_main: ProgressStyle,
89
90    #[fieldx(lazy, private, get(clone))]
91    progress_style_minor: ProgressStyle,
92
93    #[fieldx(lazy, private, get(clone))]
94    progress_style_message: ProgressStyle,
95}
96
97impl ProgressUI {
98    fn build_user_attended(&self) -> bool {
99        !self.quiet && console::user_attended()
100    }
101
102    fn build_multi_progress(&self) -> Option<MultiProgress> {
103        if self.user_attended() {
104            Some(MultiProgress::new())
105        }
106        else {
107            None
108        }
109    }
110
111    #[inline(always)]
112    fn build_progress_style_main(&self) -> ProgressStyle {
113        ProgressStyle::default_bar()
114            .template("{prefix}: [{elapsed_precise:.cyan}] [{eta:>4}] {percent_precise:>7}% {bar:30.cyan.on_240} {pos:>2.cyan}/{len:>2.cyan} {per_sec_short} {msg:.cyan}")
115            .expect("Main progress style")
116            .with_key("per_sec_short", PerSecFmt)
117            .progress_chars("█▉▊▋▌▍▎▏ ")
118    }
119
120    #[inline(always)]
121    fn build_progress_style_minor(&self) -> ProgressStyle {
122        ProgressStyle::default_bar()
123            .template("[{elapsed_precise}] {bar:5.250.on_240} {pos:>2}/{len:>2} ({prefix}) {msg}")
124            .expect("Minor progress style")
125            .progress_chars("█▉▊▋▌▍▎▏ ")
126    }
127
128    #[inline(always)]
129    fn build_progress_style_message(&self) -> ProgressStyle {
130        ProgressStyle::default_bar()
131            .template("{prefix}: {msg}")
132            .expect("Message progress style")
133    }
134
135    pub fn message_style(&self, msg_type: MsgType) -> Style {
136        match msg_type {
137            MsgType::Debug => Style::new().magenta().for_stderr(),
138            MsgType::Info => Style::new(),
139            MsgType::Warn => Style::new().yellow().for_stderr(),
140            MsgType::Error => Style::new().red().for_stderr(),
141        }
142    }
143
144    fn _println(&self, msg_type: MsgType, msg: String) {
145        if self.quiet {
146            return;
147        }
148
149        let prefix = self.message_style(msg_type).apply_to(format!("[{msg_type}]"));
150        if matches!(msg_type, MsgType::Info) {
151            println!("{prefix} {msg}");
152        }
153        else {
154            eprintln!("{prefix} {msg}");
155        }
156    }
157
158    pub fn print_message<S: ToString>(&self, msg_type: MsgType, msg: S) {
159        let msg = msg.to_string();
160        if let Some(mp) = self.multi_progress().as_ref() {
161            mp.suspend(|| {
162                self._println(msg_type, msg);
163            })
164        }
165        else {
166            self._println(msg_type, msg);
167        }
168    }
169
170    pub fn report_error<S: ToString>(&self, msg: S) {
171        self.print_message(MsgType::Error, msg);
172    }
173
174    pub fn report_warn<S: ToString>(&self, msg: S) {
175        self.print_message(MsgType::Warn, msg);
176    }
177
178    pub fn report_info<S: ToString>(&self, msg: S) {
179        self.print_message(MsgType::Info, msg);
180    }
181
182    pub fn report_debug<S: ToString>(&self, msg: S) {
183        self.print_message(MsgType::Debug, msg);
184    }
185
186    pub fn progress_style(&self, style: PStyle) -> ProgressStyle {
187        match style {
188            PStyle::Main => self.progress_style_main(),
189            PStyle::Minor => self.progress_style_minor(),
190            PStyle::Message => self.progress_style_message(),
191            PStyle::Custom(custom) => custom,
192        }
193    }
194
195    pub fn acquire_progress(&self, style: PStyle, order: Option<POrder>) -> Option<ProgressBar> {
196        let mp = self.multi_progress();
197
198        if mp.is_none() {
199            return None;
200        }
201
202        let style = self.progress_style(style);
203
204        if let Some(order) = order {
205            match order {
206                POrder::Before(before) => {
207                    mp.maybe_insert_before(before, ProgressBar::new_spinner().with_style(style.clone()))
208                }
209                POrder::After(after) => {
210                    mp.maybe_insert_after(after, ProgressBar::new_spinner().with_style(style.clone()))
211                }
212            }
213        }
214        else {
215            mp.maybe_add(ProgressBar::new(0).with_style(style.clone()))
216        }
217    }
218
219    pub fn remove(&self, pb: Option<ProgressBar>) {
220        self.multi_progress().maybe_remove(pb);
221    }
222
223    pub fn finish(&self) {
224        // self.multi_progress.clear();
225    }
226}
227
228impl Drop for ProgressUI {
229    fn drop(&mut self) {
230        self.finish();
231    }
232}