separate out the guards from the main structs

7 files changed, 133 insertions, 66 deletions

diff --git a/examples/async.rs b/examples/async.rs
index 3736cb8..b97ad24 100644
--- a/examples/async.rs
+++ b/examples/async.rs
@@ -21,9 +21,9 @@ async fn run(tm: &mut textmode::r#async::Output) -> std::io::Result<()> {
 
 fn main() {
     smol::block_on(async {
-        let mut tm = textmode::r#async::Output::new().await.unwrap();
+        let (mut tm, _guard) =
+            textmode::r#async::Output::new().await.unwrap();
         let e = run(&mut tm).await;
-        tm.cleanup().await.unwrap();
         e.unwrap();
     });
 }
diff --git a/examples/basic.rs b/examples/basic.rs
index fb13d89..f3f846d 100644
--- a/examples/basic.rs
+++ b/examples/basic.rs
@@ -1,7 +1,7 @@
 use textmode::Textmode as _;
 
 fn main() {
-    let mut tm = textmode::blocking::Output::new().unwrap();
+    let (mut tm, _guard) = textmode::blocking::Output::new().unwrap();
 
     tm.move_to(5, 5);
     tm.write_str("foo");
diff --git a/examples/tmux.rs b/examples/tmux.rs
index e5e4fcd..b85b603 100644
--- a/examples/tmux.rs
+++ b/examples/tmux.rs
@@ -310,27 +310,37 @@ impl State {
 #[must_use]
 struct Tmux {
     input: textmode::Input,
+    _raw: textmode::RawGuard,
     tm: textmode::Output,
+    _screen: textmode::ScreenGuard,
     state: State,
 }
 
 impl Tmux {
     async fn new() -> Self {
-        let input = textmode::Input::new();
-        let tm = textmode::Output::new().await.unwrap();
+        let (input, _raw) = textmode::Input::new();
+        let (tm, _screen) = textmode::Output::new().await.unwrap();
         let state = State::new();
-        Self { input, tm, state }
+        Self {
+            input,
+            _raw,
+            tm,
+            _screen,
+            state,
+        }
     }
 
     async fn run(self, ex: &smol::Executor<'_>) {
         let Self {
-            mut input,
+            input,
+            _raw,
             mut tm,
+            _screen,
             mut state,
         } = self;
 
         state.new_window(ex, state.wevents.clone());
-        state.spawn_input_task(ex, input.clone());
+        state.spawn_input_task(ex, input);
 
         ex.run(async {
             loop {
@@ -391,9 +401,6 @@ impl Tmux {
             }
         })
         .await;
-
-        tm.cleanup().await.unwrap();
-        input.cleanup();
     }
 }
 
diff --git a/src/async.rs b/src/async.rs
index 0a858c8..9e8c220 100644
--- a/src/async.rs
+++ b/src/async.rs
@@ -2,6 +2,28 @@ use futures_lite::io::AsyncWriteExt as _;
 
 use super::private::TextmodeImpl as _;
 
+pub struct ScreenGuard {
+    cleaned_up: bool,
+}
+
+impl ScreenGuard {
+    pub async fn cleanup(&mut self) -> std::io::Result<()> {
+        if self.cleaned_up {
+            return Ok(());
+        }
+        self.cleaned_up = true;
+        write_stdout(super::DEINIT).await
+    }
+}
+
+impl Drop for ScreenGuard {
+    fn drop(&mut self) {
+        futures_lite::future::block_on(async {
+            let _ = self.cleanup().await;
+        });
+    }
+}
+
 pub struct Output {
     cur: vt100::Parser,
     next: vt100::Parser,
@@ -28,7 +50,16 @@ impl super::private::TextmodeImpl for Output {
 impl super::Textmode for Output {}
 
 impl Output {
-    pub async fn new() -> std::io::Result<Self> {
+    pub async fn new() -> std::io::Result<(Self, ScreenGuard)> {
+        write_stdout(super::INIT).await?;
+
+        Ok((
+            Self::new_without_screen(),
+            ScreenGuard { cleaned_up: false },
+        ))
+    }
+
+    pub fn new_without_screen() -> Self {
         let (rows, cols) = match terminal_size::terminal_size() {
             Some((terminal_size::Width(w), terminal_size::Height(h))) => {
                 (h, w)
@@ -37,16 +68,7 @@ impl Output {
         };
         let cur = vt100::Parser::new(rows, cols, 0);
         let next = vt100::Parser::new(rows, cols, 0);
-
-        let self_ = Self { cur, next };
-        self_.write_stdout(super::INIT).await?;
-        Ok(self_)
-    }
-
-    // TODO: without async drop or async closures, i'm not sure how to do
-    // better than this
-    pub async fn cleanup(&mut self) -> std::io::Result<()> {
-        self.write_stdout(super::DEINIT).await
+        Self { cur, next }
     }
 
     pub async fn refresh(&mut self) -> std::io::Result<()> {
@@ -57,16 +79,16 @@ impl Output {
             self.next().screen().bells_diff(self.cur().screen()),
         ];
         for diff in diffs {
-            self.write_stdout(&diff).await?;
+            write_stdout(&diff).await?;
             self.cur_mut().process(&diff);
         }
         Ok(())
     }
+}
 
-    async fn write_stdout(&self, buf: &[u8]) -> std::io::Result<()> {
-        let mut stdout = blocking::Unblock::new(std::io::stdout());
-        stdout.write_all(buf).await?;
-        stdout.flush().await?;
-        Ok(())
-    }
+async fn write_stdout(buf: &[u8]) -> std::io::Result<()> {
+    let mut stdout = blocking::Unblock::new(std::io::stdout());
+    stdout.write_all(buf).await?;
+    stdout.flush().await?;
+    Ok(())
 }
diff --git a/src/blocking.rs b/src/blocking.rs
index 53cc54c..fe33e6e 100644
--- a/src/blocking.rs
+++ b/src/blocking.rs
@@ -2,6 +2,26 @@ use std::io::Write as _;
 
 use super::private::TextmodeImpl as _;
 
+pub struct ScreenGuard {
+    cleaned_up: bool,
+}
+
+impl ScreenGuard {
+    pub fn cleanup(&mut self) -> std::io::Result<()> {
+        if self.cleaned_up {
+            return Ok(());
+        }
+        self.cleaned_up = true;
+        write_stdout(super::DEINIT)
+    }
+}
+
+impl Drop for ScreenGuard {
+    fn drop(&mut self) {
+        let _ = self.cleanup();
+    }
+}
+
 pub struct Output {
     cur: vt100::Parser,
     next: vt100::Parser,
@@ -28,7 +48,15 @@ impl super::private::TextmodeImpl for Output {
 impl super::Textmode for Output {}
 
 impl Output {
-    pub fn new() -> std::io::Result<Self> {
+    pub fn new() -> std::io::Result<(Self, ScreenGuard)> {
+        write_stdout(super::INIT)?;
+        Ok((
+            Self::new_without_screen(),
+            ScreenGuard { cleaned_up: false },
+        ))
+    }
+
+    pub fn new_without_screen() -> Self {
         let (rows, cols) = match terminal_size::terminal_size() {
             Some((terminal_size::Width(w), terminal_size::Height(h))) => {
                 (h, w)
@@ -38,9 +66,7 @@ impl Output {
         let cur = vt100::Parser::new(rows, cols, 0);
         let next = vt100::Parser::new(rows, cols, 0);
 
-        let self_ = Self { cur, next };
-        self_.write_stdout(super::INIT)?;
-        Ok(self_)
+        Self { cur, next }
     }
 
     pub fn refresh(&mut self) -> std::io::Result<()> {
@@ -51,22 +77,16 @@ impl Output {
             self.next().screen().bells_diff(self.cur().screen()),
         ];
         for diff in diffs {
-            self.write_stdout(&diff)?;
+            write_stdout(&diff)?;
             self.cur_mut().process(&diff);
         }
         Ok(())
     }
-
-    fn write_stdout(&self, buf: &[u8]) -> std::io::Result<()> {
-        let mut stdout = std::io::stdout();
-        stdout.write_all(buf)?;
-        stdout.flush()?;
-        Ok(())
-    }
 }
 
-impl Drop for Output {
-    fn drop(&mut self) {
-        let _ = self.write_stdout(super::DEINIT);
-    }
+fn write_stdout(buf: &[u8]) -> std::io::Result<()> {
+    let mut stdout = std::io::stdout();
+    stdout.write_all(buf)?;
+    stdout.flush()?;
+    Ok(())
 }
diff --git a/src/input.rs b/src/input.rs
index fb0d80f..a61afa5 100644
--- a/src/input.rs
+++ b/src/input.rs
@@ -79,15 +79,39 @@ impl Key {
     }
 }
 
-#[derive(Clone)]
-pub struct Input {
+pub struct RawGuard {
     termios: nix::sys::termios::Termios,
+    cleaned_up: bool,
+}
+
+impl RawGuard {
+    pub fn cleanup(&mut self) {
+        if self.cleaned_up {
+            return;
+        }
+        self.cleaned_up = true;
+        let stdin = std::io::stdin().as_raw_fd();
+        let _ = nix::sys::termios::tcsetattr(
+            stdin,
+            nix::sys::termios::SetArg::TCSANOW,
+            &self.termios,
+        );
+    }
+}
+
+impl Drop for RawGuard {
+    fn drop(&mut self) {
+        self.cleanup();
+    }
+}
+
+pub struct Input {
     buf: Vec<u8>,
 }
 
 #[allow(clippy::new_without_default)]
 impl Input {
-    pub fn new() -> Self {
+    pub fn new() -> (Self, RawGuard) {
         let stdin = std::io::stdin().as_raw_fd();
         let termios = nix::sys::termios::tcgetattr(stdin).unwrap();
         let mut termios_raw = termios.clone();
@@ -98,8 +122,17 @@ impl Input {
             &termios_raw,
         )
         .unwrap();
+        (
+            Self::new_without_raw(),
+            RawGuard {
+                termios,
+                cleaned_up: false,
+            },
+        )
+    }
+
+    pub fn new_without_raw() -> Self {
         Self {
-            termios,
             buf: Vec::with_capacity(4096),
         }
     }
@@ -120,15 +153,6 @@ impl Input {
         self.real_read_key(false, true)
     }
 
-    pub fn cleanup(&mut self) {
-        let stdin = std::io::stdin().as_raw_fd();
-        let _ = nix::sys::termios::tcsetattr(
-            stdin,
-            nix::sys::termios::SetArg::TCSANOW,
-            &self.termios,
-        );
-    }
-
     fn real_read_key(
         &mut self,
         combine: bool,
@@ -389,9 +413,3 @@ impl std::io::Read for Input {
         std::io::stdin().read(buf)
     }
 }
-
-impl Drop for Input {
-    fn drop(&mut self) {
-        self.cleanup();
-    }
-}
diff --git a/src/lib.rs b/src/lib.rs
index 7e1842b..e0bd91f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -4,12 +4,12 @@ pub mod color;
 
 pub mod blocking;
 mod input;
-pub use input::{Input, Key};
+pub use input::{Input, Key, RawGuard};
 
 #[cfg(feature = "async")]
 pub mod r#async;
 #[cfg(feature = "async")]
-pub use r#async::Output;
+pub use r#async::{Output, ScreenGuard};
 
 const INIT: &[u8] = b"\x1b7\x1b[?47h\x1b[2J\x1b[H\x1b[?25h";
 const DEINIT: &[u8] = b"\x1b[?47l\x1b8\x1b[?25h";