Skip to main content

pito_tui/
term.rs

1use std::io::{self, Write};
2use std::panic::{self, PanicHookInfo};
3use std::sync::Arc;
4use std::sync::atomic::{AtomicU8, Ordering};
5use std::thread::{self, ThreadId};
6
7use crossterm::cursor::Show;
8use crossterm::event::{
9    DisableBracketedPaste, DisableFocusChange, DisableMouseCapture, EnableBracketedPaste,
10    EnableFocusChange, EnableMouseCapture,
11};
12use crossterm::terminal::{
13    self, BeginSynchronizedUpdate, Clear, ClearType, EndSynchronizedUpdate, EnterAlternateScreen,
14    LeaveAlternateScreen,
15};
16use crossterm::{execute, queue};
17use ratatui::backend::CrosstermBackend;
18use ratatui::{DefaultTerminal, Frame, Terminal};
19
20pub(crate) const RAW: u8 = 1;
21pub(crate) const ALTERNATE: u8 = 2;
22pub(crate) const PASTE: u8 = 4;
23pub(crate) const MOUSE: u8 = 8;
24pub(crate) const FOCUS: u8 = 16;
25
26static ON: AtomicU8 = AtomicU8::new(0);
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29#[non_exhaustive]
30pub struct Modes {
31    pub alternate: bool,
32    pub paste: bool,
33    pub mouse: bool,
34    pub focus: bool,
35}
36
37impl Modes {
38    pub const fn new() -> Self {
39        Modes {
40            alternate: true,
41            paste: true,
42            mouse: false,
43            focus: false,
44        }
45    }
46
47    pub const fn alternate(mut self, on: bool) -> Self {
48        self.alternate = on;
49        self
50    }
51
52    pub const fn paste(mut self, on: bool) -> Self {
53        self.paste = on;
54        self
55    }
56
57    pub const fn mouse(mut self, on: bool) -> Self {
58        self.mouse = on;
59        self
60    }
61
62    pub const fn focus(mut self, on: bool) -> Self {
63        self.focus = on;
64        self
65    }
66}
67
68impl Default for Modes {
69    fn default() -> Self {
70        Modes::new()
71    }
72}
73
74type Hook = dyn Fn(&PanicHookInfo<'_>) + Sync + Send + 'static;
75
76pub(crate) fn restores(ui: ThreadId, here: ThreadId, on: u8) -> bool {
77    ui == here && on & RAW != 0
78}
79
80pub(crate) fn undo(on: &AtomicU8, out: &mut impl Write, raw: impl FnOnce()) -> u8 {
81    let was = on.swap(0, Ordering::SeqCst);
82    if was & FOCUS != 0 {
83        let _ = queue!(out, DisableFocusChange);
84    }
85    if was & MOUSE != 0 {
86        let _ = queue!(out, DisableMouseCapture);
87    }
88    if was & PASTE != 0 {
89        let _ = queue!(out, DisableBracketedPaste);
90    }
91    if was & ALTERNATE != 0 {
92        let _ = queue!(out, LeaveAlternateScreen);
93    }
94    if was != 0 {
95        let _ = queue!(out, Show);
96        let _ = out.flush();
97    }
98    if was & RAW != 0 {
99        raw();
100    }
101    was
102}
103
104pub fn restore() -> bool {
105    let mut out = io::stdout().lock();
106    undo(&ON, &mut out, || {
107        let _ = terminal::disable_raw_mode();
108    }) != 0
109}
110
111fn switch(on: &AtomicU8, flag: u8, out: &mut impl Write, enable: bool) -> io::Result<()> {
112    if enable == (on.load(Ordering::SeqCst) & flag != 0) {
113        return Ok(());
114    }
115    match (flag, enable) {
116        (MOUSE, true) => execute!(out, EnableMouseCapture)?,
117        (MOUSE, false) => execute!(out, DisableMouseCapture)?,
118        (PASTE, true) => execute!(out, EnableBracketedPaste)?,
119        (PASTE, false) => execute!(out, DisableBracketedPaste)?,
120        (FOCUS, true) => execute!(out, EnableFocusChange)?,
121        (FOCUS, false) => execute!(out, DisableFocusChange)?,
122        (ALTERNATE, true) => execute!(out, EnterAlternateScreen)?,
123        (ALTERNATE, false) => execute!(out, LeaveAlternateScreen)?,
124        _ => {}
125    }
126    if enable {
127        on.fetch_or(flag, Ordering::SeqCst);
128    } else {
129        on.fetch_and(!flag, Ordering::SeqCst);
130    }
131    Ok(())
132}
133
134fn turn_on(modes: Modes) -> io::Result<()> {
135    let mut out = io::stdout().lock();
136    terminal::enable_raw_mode()?;
137    ON.fetch_or(RAW, Ordering::SeqCst);
138    switch(&ON, ALTERNATE, &mut out, modes.alternate)?;
139    switch(&ON, PASTE, &mut out, modes.paste)?;
140    switch(&ON, MOUSE, &mut out, modes.mouse)?;
141    switch(&ON, FOCUS, &mut out, modes.focus)
142}
143
144fn live() -> bool {
145    ON.load(Ordering::SeqCst) & RAW != 0
146}
147
148struct Restore {
149    previous: Arc<Hook>,
150}
151
152impl Restore {
153    fn install() -> Self {
154        let previous: Arc<Hook> = panic::take_hook().into();
155        let chained = Arc::clone(&previous);
156        let ui = thread::current().id();
157        panic::set_hook(Box::new(move |info| {
158            if restores(ui, thread::current().id(), ON.load(Ordering::SeqCst)) {
159                restore();
160            }
161            chained(info);
162        }));
163        Restore { previous }
164    }
165}
166
167impl Drop for Restore {
168    fn drop(&mut self) {
169        restore();
170        if !thread::panicking() {
171            let _ = panic::take_hook();
172            let previous = Arc::clone(&self.previous);
173            panic::set_hook(Box::new(move |info| previous(info)));
174        }
175    }
176}
177
178pub struct Term {
179    terminal: DefaultTerminal,
180    modes: Modes,
181    _restore: Restore,
182}
183
184impl Term {
185    pub fn enter(modes: Modes) -> io::Result<Term> {
186        let restore = Restore::install();
187        turn_on(modes)?;
188        execute!(io::stdout(), Clear(ClearType::All))?;
189        let terminal = Terminal::new(CrosstermBackend::new(io::stdout()))?;
190        Ok(Term {
191            terminal,
192            modes,
193            _restore: restore,
194        })
195    }
196
197    pub fn terminal(&mut self) -> &mut DefaultTerminal {
198        &mut self.terminal
199    }
200
201    pub fn modes(&self) -> Modes {
202        self.modes
203    }
204
205    pub fn draw(&mut self, render: impl FnOnce(&mut Frame)) -> io::Result<()> {
206        let mut out = io::stdout().lock();
207        if !live() {
208            return Ok(());
209        }
210        execute!(out, BeginSynchronizedUpdate)?;
211        let drawn = self.terminal.draw(render).map(|_| ());
212        execute!(out, EndSynchronizedUpdate)?;
213        drawn
214    }
215
216    pub fn mouse(&mut self, on: bool) -> io::Result<()> {
217        self.modes.mouse = on;
218        let mut out = io::stdout().lock();
219        if !live() {
220            return Ok(());
221        }
222        switch(&ON, MOUSE, &mut out, on)
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    #[test]
231    fn the_terminal_is_restored_on_the_ui_thread_only_while_raw_mode_is_on() {
232        let ui = thread::current().id();
233        let worker = thread::spawn(|| thread::current().id()).join().unwrap();
234        assert!(restores(ui, ui, RAW | ALTERNATE));
235        assert!(!restores(ui, worker, RAW | ALTERNATE));
236        assert!(!restores(ui, ui, 0));
237        assert!(!restores(ui, ui, ALTERNATE));
238    }
239
240    #[test]
241    fn every_mode_that_was_turned_on_is_turned_off_once() {
242        let on = AtomicU8::new(RAW | ALTERNATE | PASTE | MOUSE);
243        let mut out = Vec::new();
244        let mut raw = 0;
245        assert_eq!(
246            undo(&on, &mut out, || raw += 1),
247            RAW | ALTERNATE | PASTE | MOUSE
248        );
249        let written = String::from_utf8(out).unwrap();
250        let mut offs = vec!["\x1b[?2004l", "\x1b[?1049l", "\x1b[?25h"];
251        if cfg!(not(windows)) {
252            offs.push("\x1b[?1000l");
253        }
254        for off in offs {
255            assert!(written.contains(off), "{off:?} in {written:?}");
256        }
257        assert!(!written.contains("\x1b[?1004l"), "focus was never on");
258        assert_eq!(raw, 1);
259        let mut again = Vec::new();
260        assert_eq!(undo(&on, &mut again, || raw += 1), 0);
261        assert!(again.is_empty());
262        assert_eq!(raw, 1);
263    }
264
265    #[test]
266    fn any_thread_gives_the_terminal_back_and_only_the_first_one_does() {
267        ON.store(RAW | PASTE, Ordering::SeqCst);
268        let threads: Vec<_> = (0..4).map(|_| thread::spawn(restore)).collect();
269        let gave = threads
270            .into_iter()
271            .filter_map(|thread| thread.join().ok())
272            .filter(|gave| *gave)
273            .count();
274        assert_eq!(gave, 1);
275        assert_eq!(ON.load(Ordering::SeqCst), 0);
276        assert!(!restore());
277    }
278}