use std::cell::RefCell;
use std::io::{IsTerminal, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, Weak};
use std::time::Instant;
const FRAMES: [char; 10] = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
const MAX_LINES: usize = 12;
struct Slot {
name: String,
started: Instant,
buf: Vec<u8>,
running: bool,
done: bool,
}
pub struct Stage {
slots: Mutex<Vec<Slot>>,
out: Mutex<usize>,
live: bool,
stop: AtomicBool,
}
thread_local! {
static SINK: RefCell<Option<(Arc<Stage>, usize)>> = const { RefCell::new(None) };
}
impl Stage {
pub fn begin(names: &[&str]) -> Arc<Stage> {
let now = Instant::now();
let stage = Arc::new(Stage {
slots: Mutex::new(
names
.iter()
.map(|n| Slot {
name: crate::ui::sanitize(
n.strip_prefix("pre-commit-")
.or_else(|| n.strip_prefix("pre-push-"))
.unwrap_or(n),
),
started: now,
buf: Vec::new(),
running: false,
done: false,
})
.collect(),
),
out: Mutex::new(0),
live: enabled() && watching(),
stop: AtomicBool::new(false),
});
if stage.live {
let weak = Arc::downgrade(&stage);
let _ = std::thread::Builder::new()
.name("amont-live".into())
.spawn(move || tick(weak));
}
stage
}
pub fn enter(self: &Arc<Stage>, idx: usize) -> SinkGuard {
{
let mut slots = self.slots.lock().unwrap_or_else(|p| p.into_inner());
if let Some(slot) = slots.get_mut(idx) {
slot.running = true;
slot.started = Instant::now();
}
}
SINK.with(|s| *s.borrow_mut() = Some((Arc::clone(self), idx)));
SinkGuard
}
pub fn append_raw(&self, idx: usize, bytes: &[u8]) {
let mut slots = self.slots.lock().unwrap_or_else(|p| p.into_inner());
if let Some(slot) = slots.get_mut(idx) {
if !slot.done {
slot.buf.extend_from_slice(bytes);
}
}
}
fn append_line(&self, idx: usize, line: &str) {
let mut slots = self.slots.lock().unwrap_or_else(|p| p.into_inner());
if let Some(slot) = slots.get_mut(idx) {
if !slot.done {
slot.buf.extend_from_slice(line.as_bytes());
slot.buf.push(b'\n');
}
}
}
pub fn finish(&self, idx: usize) {
let block = {
let mut slots = self.slots.lock().unwrap_or_else(|p| p.into_inner());
let Some(slot) = slots.get_mut(idx) else {
return;
};
slot.done = true;
slot.running = false;
std::mem::take(&mut slot.buf)
};
if block.is_empty() && !self.live {
return;
}
let mut drawn = self.out.lock().unwrap_or_else(|p| p.into_inner());
if !block.is_empty() {
if *drawn > 0 {
let mut err = std::io::stderr().lock();
let _ = write!(err, "\x1b[{}A\x1b[J", *drawn);
let _ = err.flush();
*drawn = 0;
}
let stdout = std::io::stdout();
let mut handle = stdout.lock();
let _ = handle.write_all(&block);
let _ = handle.flush();
}
self.repaint(&mut drawn);
}
fn repaint(&self, drawn: &mut usize) {
if !self.live {
return;
}
let entries: Vec<(String, f64)> = {
let slots = self.slots.lock().unwrap_or_else(|p| p.into_inner());
let now = Instant::now();
slots
.iter()
.filter(|s| s.running && !s.done)
.map(|s| (s.name.clone(), now.duration_since(s.started).as_secs_f64()))
.collect()
};
let text = region(&entries, term_width());
let mut paint = String::new();
if *drawn > 0 {
paint.push_str(&format!("\x1b[{}A\x1b[J", *drawn));
}
paint.push_str(&text);
if paint.is_empty() {
return;
}
let mut err = std::io::stderr().lock();
let _ = err.write_all(paint.as_bytes());
let _ = err.flush();
*drawn = text.matches('\n').count();
}
}
impl Drop for Stage {
fn drop(&mut self) {
self.stop.store(true, Ordering::Relaxed);
if !self.live {
return;
}
let drawn = self.out.get_mut().unwrap_or_else(|p| p.into_inner());
if *drawn > 0 {
let mut err = std::io::stderr().lock();
let _ = write!(err, "\x1b[{}A\x1b[J", *drawn);
let _ = err.flush();
*drawn = 0;
}
}
}
fn tick(weak: Weak<Stage>) {
loop {
std::thread::sleep(std::time::Duration::from_millis(80));
let Some(stage) = weak.upgrade() else { return };
if stage.stop.load(Ordering::Relaxed) {
return;
}
let mut drawn = stage.out.lock().unwrap_or_else(|p| p.into_inner());
stage.repaint(&mut drawn);
}
}
fn region(entries: &[(String, f64)], width: usize) -> String {
if entries.is_empty() {
return String::new();
}
let pad = entries
.iter()
.take(MAX_LINES)
.map(|(name, _)| name.chars().count())
.max()
.unwrap_or(0);
let mut out = String::new();
for (name, secs) in entries.iter().take(MAX_LINES) {
let frame = FRAMES[((secs * 10.0) as usize) % FRAMES.len()];
let line = format!("{frame} {name:<pad$} {secs:>5.1}s");
if line.chars().count() > width {
out.extend(line.chars().take(width));
} else {
out.push_str(&line);
}
out.push('\n');
}
if entries.len() > MAX_LINES {
out.push_str(&format!("… and {} more\n", entries.len() - MAX_LINES));
}
out
}
fn term_width() -> usize {
std::env::var("COLUMNS")
.ok()
.and_then(|c| c.parse::<usize>().ok())
.filter(|w| *w >= 20)
.unwrap_or(100)
}
pub struct FinishOnDrop<'a> {
stage: &'a Stage,
idx: usize,
}
impl<'a> FinishOnDrop<'a> {
pub fn new(stage: &'a Stage, idx: usize) -> FinishOnDrop<'a> {
FinishOnDrop { stage, idx }
}
}
impl Drop for FinishOnDrop<'_> {
fn drop(&mut self) {
self.stage.finish(self.idx);
}
}
pub struct SinkGuard;
impl Drop for SinkGuard {
fn drop(&mut self) {
SINK.with(|s| *s.borrow_mut() = None);
}
}
pub fn current_sink() -> Option<(Arc<Stage>, usize)> {
SINK.with(|s| s.borrow().clone())
}
pub fn say(line: &str) {
let routed = SINK.with(|s| {
s.borrow().as_ref().map(|(stage, idx)| {
stage.append_line(*idx, line);
})
});
if routed.is_none() {
println!("{line}");
}
}
#[macro_export]
macro_rules! say {
($($arg:tt)*) => {
$crate::live::say(&format!($($arg)*))
};
}
pub fn enabled() -> bool {
static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
*ENABLED.get_or_init(|| crate::config::boolean_or("amont.progress", true))
}
pub fn watching() -> bool {
static WATCHING: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
*WATCHING.get_or_init(|| {
if !std::io::stderr().is_terminal() {
return false;
}
match std::env::var("TERM") {
Ok(term) => term != "dumb",
Err(_) => !cfg!(windows),
}
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn slots_do_not_share_a_buffer() {
let stage = Stage::begin(&["a", "b"]);
std::thread::scope(|scope| {
for idx in 0..2 {
let stage = Arc::clone(&stage);
scope.spawn(move || {
let _guard = stage.enter(idx);
for i in 0..50 {
say(&format!("check-{idx} line-{i}"));
std::thread::yield_now();
}
});
}
});
let slots = stage.slots.lock().unwrap();
for idx in 0..2 {
let text = String::from_utf8(slots[idx].buf.clone()).unwrap();
assert_eq!(text.lines().count(), 50);
assert!(
text.lines()
.all(|l| l.starts_with(&format!("check-{idx} "))),
"a foreign line landed in slot {idx}"
);
}
}
#[test]
fn no_sink_means_no_capture() {
let stage = Stage::begin(&["a"]);
say("goes to stdout, not to a slot");
let slots = stage.slots.lock().unwrap();
assert!(slots[0].buf.is_empty());
}
#[test]
fn a_finished_slot_takes_no_more_writes() {
let stage = Stage::begin(&["a"]);
stage.append_raw(0, b"before\n");
stage.finish(0);
stage.append_raw(0, b"after\n");
let slots = stage.slots.lock().unwrap();
assert!(slots[0].buf.is_empty(), "a write landed after finish");
}
#[test]
fn a_slot_name_is_sanitised_at_begin() {
let stage = Stage::begin(&["evil\u{1b}[2Jname\rhere"]);
let slots = stage.slots.lock().unwrap();
assert!(!slots[0].name.contains('\u{1b}'), "{:?}", slots[0].name);
assert!(!slots[0].name.contains('\r'), "{:?}", slots[0].name);
}
#[test]
fn a_slot_name_drops_the_stage_prefix() {
let stage = Stage::begin(&["pre-commit-clippy", "pre-push-run-tests", "bare"]);
let slots = stage.slots.lock().unwrap();
assert_eq!(slots[0].name, "clippy");
assert_eq!(slots[1].name, "run-tests");
assert_eq!(slots[2].name, "bare");
}
#[test]
fn frames_advance_with_time() {
let a = region(&[("clippy".into(), 0.0)], 80);
let b = region(&[("clippy".into(), 0.1)], 80);
let c = region(&[("clippy".into(), 1.0)], 80);
assert_ne!(a.chars().next(), b.chars().next());
assert_eq!(a.chars().next(), c.chars().next(), "10 frames per second");
}
#[test]
fn region_lines_align() {
let text = region(&[("a".into(), 0.0), ("longer-name".into(), 0.0)], 80);
let widths: Vec<usize> = text.lines().map(|l| l.chars().count()).collect();
assert_eq!(widths[0], widths[1], "{text:?}");
}
#[test]
fn region_caps_and_counts_the_rest() {
let entries: Vec<(String, f64)> = (0..13).map(|i| (format!("check-{i}"), 0.0)).collect();
let text = region(&entries, 80);
assert_eq!(text.lines().count(), MAX_LINES + 1);
assert!(text.ends_with("… and 1 more\n"), "{text:?}");
}
#[test]
fn region_respects_width() {
let text = region(&[("a-name-much-longer-than-the-terminal".into(), 0.0)], 20);
assert!(text.lines().all(|l| l.chars().count() <= 20), "{text:?}");
}
#[test]
fn an_empty_region_is_empty() {
assert_eq!(region(&[], 80), "");
}
}