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}