1use std::{
2 borrow::Cow,
3 collections::HashMap,
4 io::{IsTerminal, Write},
5 sync::{Arc, Mutex},
6 time::{Duration, Instant},
7};
8
9pub struct Progress {
10 inner: Mutex<ProgressInner>,
11}
12
13struct ProgressInner {
14 completed: Vec<CompletedTask>,
15 in_progress: HashMap<String, InProgressTask>,
16 is_tty: bool,
17 terminal_width: usize,
18 max_example_name_len: usize,
19 status_lines_displayed: usize,
20 total_examples: usize,
21}
22
23#[derive(Clone)]
24struct InProgressTask {
25 step: Step,
26 started_at: Instant,
27 substatus: Option<String>,
28 info: Option<(String, InfoStyle)>,
29}
30
31struct CompletedTask {
32 step: Step,
33 example_name: String,
34 duration: Duration,
35 info: Option<(String, InfoStyle)>,
36}
37
38#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
39pub enum Step {
40 LoadProject,
41 Context,
42 FormatPrompt,
43 Predict,
44 Score,
45}
46
47#[derive(Clone, Copy, Debug, PartialEq, Eq)]
48pub enum InfoStyle {
49 Normal,
50 Warning,
51}
52
53impl Step {
54 pub fn label(&self) -> &'static str {
55 match self {
56 Step::LoadProject => "Load",
57 Step::Context => "Context",
58 Step::FormatPrompt => "Format",
59 Step::Predict => "Predict",
60 Step::Score => "Score",
61 }
62 }
63
64 fn color_code(&self) -> &'static str {
65 match self {
66 Step::LoadProject => "\x1b[33m",
67 Step::Context => "\x1b[35m",
68 Step::FormatPrompt => "\x1b[34m",
69 Step::Predict => "\x1b[32m",
70 Step::Score => "\x1b[31m",
71 }
72 }
73}
74
75const RIGHT_MARGIN: usize = 4;
76
77impl Progress {
78 pub fn new(total_examples: usize) -> Arc<Self> {
79 Arc::new(Self {
80 inner: Mutex::new(ProgressInner {
81 completed: Vec::new(),
82 in_progress: HashMap::new(),
83 is_tty: std::io::stderr().is_terminal(),
84 terminal_width: get_terminal_width(),
85 max_example_name_len: 0,
86 status_lines_displayed: 0,
87 total_examples,
88 }),
89 })
90 }
91
92 pub fn start(self: &Arc<Self>, step: Step, example_name: &str) -> Arc<StepProgress> {
93 {
94 let mut inner = self.inner.lock().unwrap();
95
96 Self::clear_status_lines(&mut inner);
97
98 inner.max_example_name_len = inner.max_example_name_len.max(example_name.len());
99
100 inner.in_progress.insert(
101 example_name.to_string(),
102 InProgressTask {
103 step,
104 started_at: Instant::now(),
105 substatus: None,
106 info: None,
107 },
108 );
109
110 Self::print_status_lines(&mut inner);
111 }
112
113 Arc::new(StepProgress {
114 progress: self.clone(),
115 step,
116 example_name: example_name.to_string(),
117 })
118 }
119
120 pub fn finish(&self, step: Step, example_name: &str) {
121 let mut inner = self.inner.lock().unwrap();
122
123 let task = inner.in_progress.remove(example_name);
124 if let Some(task) = task {
125 if task.step == step {
126 inner.completed.push(CompletedTask {
127 step: task.step,
128 example_name: example_name.to_string(),
129 duration: task.started_at.elapsed(),
130 info: task.info,
131 });
132
133 Self::clear_status_lines(&mut inner);
134 Self::print_completed(&inner, inner.completed.last().unwrap());
135 Self::print_status_lines(&mut inner);
136 } else {
137 inner.in_progress.insert(example_name.to_string(), task);
138 }
139 }
140 }
141
142 fn clear_status_lines(inner: &mut ProgressInner) {
143 if inner.is_tty && inner.status_lines_displayed > 0 {
144 // Move up and clear each line we previously displayed
145 for _ in 0..inner.status_lines_displayed {
146 eprint!("\x1b[A\x1b[K");
147 }
148 let _ = std::io::stderr().flush();
149 inner.status_lines_displayed = 0;
150 }
151 }
152
153 fn print_completed(inner: &ProgressInner, task: &CompletedTask) {
154 let duration = format_duration(task.duration);
155 let name_width = inner.max_example_name_len;
156
157 if inner.is_tty {
158 let reset = "\x1b[0m";
159 let bold = "\x1b[1m";
160 let dim = "\x1b[2m";
161
162 let yellow = "\x1b[33m";
163 let info_part = task
164 .info
165 .as_ref()
166 .map(|(s, style)| {
167 if *style == InfoStyle::Warning {
168 format!("{yellow}{s}{reset}")
169 } else {
170 s.to_string()
171 }
172 })
173 .unwrap_or_default();
174
175 let prefix = format!(
176 "{bold}{color}{label:>12}{reset} {name:<name_width$} {dim}│{reset} {info_part}",
177 color = task.step.color_code(),
178 label = task.step.label(),
179 name = task.example_name,
180 );
181
182 let duration_with_margin = format!("{duration} ");
183 let padding_needed = inner
184 .terminal_width
185 .saturating_sub(RIGHT_MARGIN)
186 .saturating_sub(duration_with_margin.len())
187 .saturating_sub(strip_ansi_len(&prefix));
188 let padding = " ".repeat(padding_needed);
189
190 eprintln!("{prefix}{padding}{dim}{duration_with_margin}{reset}");
191 } else {
192 let info_part = task
193 .info
194 .as_ref()
195 .map(|(s, _)| format!(" | {}", s))
196 .unwrap_or_default();
197
198 eprintln!(
199 "{label:>12} {name:<name_width$}{info_part} {duration}",
200 label = task.step.label(),
201 name = task.example_name,
202 );
203 }
204 }
205
206 fn print_status_lines(inner: &mut ProgressInner) {
207 if !inner.is_tty || inner.in_progress.is_empty() {
208 inner.status_lines_displayed = 0;
209 return;
210 }
211
212 let reset = "\x1b[0m";
213 let bold = "\x1b[1m";
214 let dim = "\x1b[2m";
215
216 // Build the done/in-progress/total label
217 let done_count = inner.completed.len();
218 let in_progress_count = inner.in_progress.len();
219 let range_label = format!(
220 " {}/{}/{} ",
221 done_count, in_progress_count, inner.total_examples
222 );
223
224 // Print a divider line with range label aligned with timestamps
225 let range_visible_len = range_label.len();
226 let left_divider_len = inner
227 .terminal_width
228 .saturating_sub(RIGHT_MARGIN)
229 .saturating_sub(range_visible_len);
230 let left_divider = "─".repeat(left_divider_len);
231 let right_divider = "─".repeat(RIGHT_MARGIN);
232 eprintln!("{dim}{left_divider}{reset}{range_label}{dim}{right_divider}{reset}");
233
234 let mut tasks: Vec<_> = inner.in_progress.iter().collect();
235 tasks.sort_by_key(|(name, _)| *name);
236
237 let mut lines_printed = 0;
238
239 for (name, task) in tasks.iter() {
240 let elapsed = format_duration(task.started_at.elapsed());
241 let substatus_part = task
242 .substatus
243 .as_ref()
244 .map(|s| truncate_with_ellipsis(s, 30))
245 .unwrap_or_default();
246
247 let step_label = task.step.label();
248 let step_color = task.step.color_code();
249 let name_width = inner.max_example_name_len;
250
251 let prefix = format!(
252 "{bold}{step_color}{step_label:>12}{reset} {name:<name_width$} {dim}│{reset} {substatus_part}",
253 name = name,
254 );
255
256 let duration_with_margin = format!("{elapsed} ");
257 let padding_needed = inner
258 .terminal_width
259 .saturating_sub(RIGHT_MARGIN)
260 .saturating_sub(duration_with_margin.len())
261 .saturating_sub(strip_ansi_len(&prefix));
262 let padding = " ".repeat(padding_needed);
263
264 eprintln!("{prefix}{padding}{dim}{duration_with_margin}{reset}");
265 lines_printed += 1;
266 }
267
268 inner.status_lines_displayed = lines_printed + 1; // +1 for the divider line
269 let _ = std::io::stderr().flush();
270 }
271
272 pub fn clear(&self) {
273 let mut inner = self.inner.lock().unwrap();
274 Self::clear_status_lines(&mut inner);
275 }
276}
277
278pub struct StepProgress {
279 progress: Arc<Progress>,
280 step: Step,
281 example_name: String,
282}
283
284impl StepProgress {
285 pub fn set_substatus(&self, substatus: impl Into<Cow<'static, str>>) {
286 let mut inner = self.progress.inner.lock().unwrap();
287 if let Some(task) = inner.in_progress.get_mut(&self.example_name) {
288 task.substatus = Some(substatus.into().into_owned());
289 Progress::clear_status_lines(&mut inner);
290 Progress::print_status_lines(&mut inner);
291 }
292 }
293
294 pub fn clear_substatus(&self) {
295 let mut inner = self.progress.inner.lock().unwrap();
296 if let Some(task) = inner.in_progress.get_mut(&self.example_name) {
297 task.substatus = None;
298 Progress::clear_status_lines(&mut inner);
299 Progress::print_status_lines(&mut inner);
300 }
301 }
302
303 pub fn set_info(&self, info: impl Into<String>, style: InfoStyle) {
304 let mut inner = self.progress.inner.lock().unwrap();
305 if let Some(task) = inner.in_progress.get_mut(&self.example_name) {
306 task.info = Some((info.into(), style));
307 }
308 }
309}
310
311impl Drop for StepProgress {
312 fn drop(&mut self) {
313 self.progress.finish(self.step, &self.example_name);
314 }
315}
316
317#[cfg(unix)]
318fn get_terminal_width() -> usize {
319 unsafe {
320 let mut winsize: libc::winsize = std::mem::zeroed();
321 if libc::ioctl(libc::STDERR_FILENO, libc::TIOCGWINSZ, &mut winsize) == 0
322 && winsize.ws_col > 0
323 {
324 winsize.ws_col as usize
325 } else {
326 80
327 }
328 }
329}
330
331#[cfg(not(unix))]
332fn get_terminal_width() -> usize {
333 80
334}
335
336fn strip_ansi_len(s: &str) -> usize {
337 let mut len = 0;
338 let mut in_escape = false;
339 for c in s.chars() {
340 if c == '\x1b' {
341 in_escape = true;
342 } else if in_escape {
343 if c == 'm' {
344 in_escape = false;
345 }
346 } else {
347 len += 1;
348 }
349 }
350 len
351}
352
353fn truncate_with_ellipsis(s: &str, max_len: usize) -> String {
354 if s.len() <= max_len {
355 s.to_string()
356 } else {
357 format!("{}…", &s[..max_len.saturating_sub(1)])
358 }
359}
360
361fn format_duration(duration: Duration) -> String {
362 const MINUTE_IN_MILLIS: f32 = 60. * 1000.;
363
364 let millis = duration.as_millis() as f32;
365 if millis < 1000.0 {
366 format!("{}ms", millis)
367 } else if millis < MINUTE_IN_MILLIS {
368 format!("{:.1}s", millis / 1_000.0)
369 } else {
370 format!("{:.1}m", millis / MINUTE_IN_MILLIS)
371 }
372}