Merge pull request #210 from solidiquis/sigint-tty

Restore tty settings

8 files changed, 163 insertions, 62 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 7d9fa5f..a879ded 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -253,6 +253,16 @@ dependencies = [
 ]
 
 [[package]]
+name = "ctrlc"
+version = "3.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2a011bbe2c35ce9c1f143b7af6f94f29a167beb4cd1d29e6740ce836f723120e"
+dependencies = [
+ "nix",
+ "windows-sys 0.48.0",
+]
+
+[[package]]
 name = "cxx"
 version = "1.0.94"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -342,6 +352,7 @@ dependencies = [
  "clap_complete",
  "config",
  "crossterm",
+ "ctrlc",
  "dirs",
  "errno 0.3.1",
  "filesize",
@@ -673,6 +684,18 @@ dependencies = [
 ]
 
 [[package]]
+name = "nix"
+version = "0.26.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bfdda3d196821d6af13126e40375cdf7da646a96114af134d5f417a9a1dc8e1a"
+dependencies = [
+ "bitflags",
+ "cfg-if",
+ "libc",
+ "static_assertions",
+]
+
+[[package]]
 name = "nom"
 version = "7.1.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1046,6 +1069,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a507befe795404456341dfab10cef66ead4c041f62b8b11bbb92bffe5d0953e0"
 
 [[package]]
+name = "static_assertions"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
+
+[[package]]
 name = "strip-ansi-escapes"
 version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index 7982fc6..07c9a16 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -30,6 +30,7 @@ clap = { version = "4.1.1", features = ["derive"] }
 clap_complete = "4.1.1"
 config = { version = "0.13.3", features = ["toml"] }
 crossterm = "0.26.1"
+ctrlc = "3.4.0"
 dirs = "5.0"
 errno = "0.3.1"
 filesize = "0.2.0"
diff --git a/src/main.rs b/src/main.rs
index 2bb1f1a..5c29e85 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -23,7 +23,7 @@ use clap::CommandFactory;
 use context::{layout, Context};
 use progress::Message;
 use render::{Engine, Flat, FlatInverted, Inverted, Regular};
-use std::{error::Error, io::stdout, process::ExitCode};
+use std::{error::Error, io::stdout, process::ExitCode, sync::Arc};
 use tree::Tree;
 
 /// Operations to wrangle ANSI escaped strings.
@@ -64,7 +64,11 @@ mod tty;
 mod utils;
 
 fn main() -> ExitCode {
-    if let Err(e) = run() {
+    let result = run();
+
+    tty::restore_tty();
+
+    if let Err(e) = result {
         eprintln!("{e}");
         return ExitCode::FAILURE;
     }
@@ -80,40 +84,53 @@ fn run() -> Result<(), Box<dyn Error>> {
     }
 
     context::color::no_color_env();
+
     styles::init(ctx.no_color());
 
-    let indicator = (ctx.stdout_is_tty && !ctx.no_progress).then(progress::Indicator::measure);
+    let indicator = (ctx.stdout_is_tty && !ctx.no_progress)
+        .then(progress::Indicator::measure)
+        .map(Arc::new);
 
-    let (tree, ctx) =
-        Tree::try_init_and_update_context(ctx, indicator.as_ref()).map_err(|err| {
-            if let Some(ref progress) = indicator {
-                progress.mailbox().send(Message::RenderReady).unwrap();
-            }
-            err
+    if indicator.is_some() {
+        let indicator = indicator.clone();
+
+        ctrlc::set_handler(move || {
+            let _ = progress::IndicatorHandle::terminate(indicator.clone());
+            tty::restore_tty();
         })?;
+    }
 
-    let output = match ctx.layout {
-        layout::Type::Flat => {
-            let render = Engine::<Flat>::new(tree, ctx);
-            format!("{render}")
+    let (tree, ctx) = match Tree::try_init(ctx, indicator.clone()) {
+        Ok(res) => res,
+        Err(err) => {
+            let _ = progress::IndicatorHandle::terminate(indicator);
+            return Err(Box::new(err));
         },
-        layout::Type::Iflat => {
-            let render = Engine::<FlatInverted>::new(tree, ctx);
-            format!("{render}")
-        },
-        layout::Type::Inverted => {
-            let render = Engine::<Inverted>::new(tree, ctx);
-            format!("{render}")
-        },
-        layout::Type::Regular => {
-            let render = Engine::<Regular>::new(tree, ctx);
+    };
+
+    macro_rules! compute_output {
+        ($t:ty) => {{
+            let render = Engine::<$t>::new(tree, ctx);
             format!("{render}")
-        },
+        }};
+    }
+
+    let output = match ctx.layout {
+        layout::Type::Flat => compute_output!(Flat),
+        layout::Type::Iflat => compute_output!(FlatInverted),
+        layout::Type::Inverted => compute_output!(Inverted),
+        layout::Type::Regular => compute_output!(Regular),
     };
 
-    if let Some(progress) = indicator {
+    if let Some(mut progress) = indicator {
         progress.mailbox().send(Message::RenderReady)?;
-        progress.join_handle.join().unwrap()?;
+
+        if let Some(hand) = Arc::get_mut(&mut progress) {
+            hand.join_handle
+                .take()
+                .map(|h| h.join().unwrap())
+                .transpose()?;
+        }
     }
 
     #[cfg(debug_assertions)]
diff --git a/src/progress.rs b/src/progress.rs
index 12b3fde..9acc332 100644
--- a/src/progress.rs
+++ b/src/progress.rs
@@ -5,10 +5,16 @@ use crossterm::{
 };
 use std::{
     io::{self, Write},
-    sync::mpsc::{self, Sender},
-    thread,
+    sync::{
+        mpsc::{self, SendError, SyncSender},
+        Arc,
+    },
+    thread::{self, JoinHandle},
+    time::Duration,
 };
 
+const PRIORITY_MAIL_TIMEOUT: Duration = Duration::from_nanos(1);
+
 /// Responsible for displying the progress indicator. This struct will be owned by a separate
 /// thread that is responsible for displaying the progress text whereas the [`IndicatorHandle`]
 /// is how the outside world will interact with it.
@@ -20,10 +26,11 @@ pub struct Indicator<'a> {
 
 /// This struct is how the outside world will inform the [`Indicator`] about the progress of the
 /// program. The `join_handle` returns the handle to the thread that owns the [`Indicator`] and the
-/// `mailbox` is the [`Sender`] channel that allows [`Message`]s to be sent to [`Indicator`].
+/// `mailbox` is the [`SyncSender`] channel that allows [`Message`]s to be sent to [`Indicator`].
 pub struct IndicatorHandle {
-    pub join_handle: thread::JoinHandle<Result<(), Error>>,
-    mailbox: Sender<Message>,
+    pub join_handle: Option<JoinHandle<Result<(), Error>>>,
+    mailbox: SyncSender<Message>,
+    priority_mailbox: SyncSender<()>,
 }
 
 /// The different messages that could be sent to the thread that owns the [`Indicator`].
@@ -41,7 +48,7 @@ pub enum Message {
 }
 
 /// All of the different states the [`Indicator`] can be in during its life cycle.
-#[derive(Default, PartialEq)]
+#[derive(Debug, Default, PartialEq)]
 enum IndicatorState {
     /// We are currently reading from disk.
     #[default]
@@ -57,7 +64,13 @@ enum IndicatorState {
 /// Errors associated with [`crossterm`];
 #[derive(Debug, thiserror::Error)]
 #[error(transparent)]
-pub struct Error(#[from] io::Error);
+pub enum Error {
+    #[error("#{0}")]
+    Io(#[from] io::Error),
+
+    #[error("#{0}")]
+    Send(#[from] SendError<()>),
+}
 
 impl Default for Indicator<'_> {
     /// Default constructor for [`Indicator`].
@@ -73,19 +86,40 @@ impl Default for Indicator<'_> {
 impl IndicatorHandle {
     /// The constructor for an [`IndicatorHandle`].
     pub fn new(
-        join_handle: thread::JoinHandle<Result<(), Error>>,
-        mailbox: Sender<Message>,
+        join_handle: Option<JoinHandle<Result<(), Error>>>,
+        mailbox: SyncSender<Message>,
+        priority_mailbox: SyncSender<()>,
     ) -> Self {
         Self {
             join_handle,
             mailbox,
+            priority_mailbox,
         }
     }
 
     /// Getter for a cloned `mailbox` wherewith to send [`Message`]s to the [`Indicator`].
-    pub fn mailbox(&self) -> Sender<Message> {
+    pub fn mailbox(&self) -> SyncSender<Message> {
         self.mailbox.clone()
     }
+
+    /// Getter for a cloned `priority_mailbox` wherewith to send [`Message`]s to the [`Indicator`].
+    pub fn priority_mailbox(&self) -> SyncSender<()> {
+        self.priority_mailbox.clone()
+    }
+
+    pub fn terminate(this: Option<Arc<Self>>) -> Result<(), Error> {
+        if let Some(mut handle) = this {
+            handle.priority_mailbox().send(())?;
+
+            if let Some(hand) = Arc::get_mut(&mut handle) {
+                hand.join_handle
+                    .take()
+                    .map(|h| h.join().unwrap())
+                    .transpose()?;
+            }
+        }
+        Ok(())
+    }
 }
 
 impl<'a> Indicator<'a> {
@@ -93,7 +127,8 @@ impl<'a> Indicator<'a> {
     /// through its internal states. An [`IndicatorHandle`] is returned as a mechanism to allow the
     /// outside world to send messages to the worker thread and ultimately to the [`Indicator`].
     pub fn measure() -> IndicatorHandle {
-        let (tx, rx) = mpsc::channel();
+        let (tx, rx) = mpsc::sync_channel(1024);
+        let (ptx, prx) = mpsc::sync_channel(1);
 
         let join_handle = thread::spawn(move || {
             let mut indicator = Self::default();
@@ -102,27 +137,29 @@ impl<'a> Indicator<'a> {
             indicator.stdout.execute(cursor::Hide)?;
 
             while let Ok(msg) = rx.recv() {
-                if indicator.state == IndicatorState::Indexing {
-                    match msg {
-                        Message::Index => indicator.index()?,
-                        Message::DoneIndexing => {
-                            indicator.update_state(IndicatorState::Rendering)?;
-                        },
-                        Message::RenderReady => {},
-                    }
-                }
-
-                if indicator.state == IndicatorState::Rendering && msg == Message::RenderReady {
+                if prx.recv_timeout(PRIORITY_MAIL_TIMEOUT).is_ok() {
                     indicator.update_state(IndicatorState::Done)?;
                     break;
                 }
+
+                match msg {
+                    Message::Index => indicator.index()?,
+                    Message::DoneIndexing => {
+                        indicator.update_state(IndicatorState::Rendering)?;
+                    },
+                    Message::RenderReady => {
+                        indicator.update_state(IndicatorState::Done)?;
+                        return Ok(());
+                    },
+                }
+
                 indicator.stdout.execute(cursor::RestorePosition)?;
             }
 
             Ok(())
         });
 
-        IndicatorHandle::new(join_handle, tx)
+        IndicatorHandle::new(Some(join_handle), tx, ptx)
     }
 
     /// Updates the `state` of the [`Indicator`] to `new_state`, immediately running an associated
@@ -139,11 +176,10 @@ impl<'a> Indicator<'a> {
                 self.rendering();
             },
 
-            (Rendering, Done) => {
+            (Rendering | Indexing, Done) => {
                 let stdout = &mut self.stdout;
                 stdout.execute(terminal::Clear(ClearType::CurrentLine))?;
                 stdout.execute(cursor::RestorePosition)?;
-                stdout.execute(cursor::Show)?;
             },
             _ => (),
         }
diff --git a/src/tree/error.rs b/src/tree/error.rs
index f95e40d..570ee79 100644
--- a/src/tree/error.rs
+++ b/src/tree/error.rs
@@ -33,8 +33,11 @@ pub enum Error {
 
     #[cfg(unix)]
     #[error("{0}")]
-    Persmissions(#[from] PermissionsError),
+    Permissions(#[from] PermissionsError),
 
     #[error("{0}")]
     UninitializedTheme(#[from] StyleError<'static>),
+
+    #[error("Terminated erdtree...")]
+    Terminated,
 }
diff --git a/src/tree/mod.rs b/src/tree/mod.rs
index f7565a6..fff0815 100644
--- a/src/tree/mod.rs
+++ b/src/tree/mod.rs
@@ -2,7 +2,7 @@ use crate::{
     context::{column, Context},
     disk_usage::file_size::FileSize,
     fs::inode::Inode,
-    progress::{self, IndicatorHandle, Message},
+    progress::{IndicatorHandle, Message},
     utils,
 };
 use count::FileCount;
@@ -16,7 +16,10 @@ use std::{
     fs,
     path::PathBuf,
     result::Result as StdResult,
-    sync::mpsc::{self, Sender},
+    sync::{
+        mpsc::{self, Sender},
+        Arc,
+    },
     thread,
 };
 use visitor::{BranchVisitorBuilder, TraversalState};
@@ -52,9 +55,9 @@ impl Tree {
 
     /// Initiates file-system traversal and [Tree] as well as updates the [Context] object with
     /// various properties necessary to render output.
-    pub fn try_init_and_update_context(
+    pub fn try_init(
         mut ctx: Context,
-        indicator: Option<&IndicatorHandle>,
+        indicator: Option<Arc<IndicatorHandle>>,
     ) -> Result<(Self, Context)> {
         let mut column_properties = column::Properties::from(&ctx);
 
@@ -102,12 +105,12 @@ impl Tree {
     fn traverse(
         ctx: &Context,
         column_properties: &mut column::Properties,
-        indicator: Option<&IndicatorHandle>,
+        indicator: Option<Arc<IndicatorHandle>>,
     ) -> Result<(Arena<Node>, NodeId)> {
         let walker = WalkParallel::try_from(ctx)?;
         let (tx, rx) = mpsc::channel();
 
-        let progress_indicator_mailbox = indicator.map(progress::IndicatorHandle::mailbox);
+        let progress_indicator_mailbox = indicator.map(|arc| arc.mailbox());
 
         thread::scope(|s| {
             let res = s.spawn(move || {
@@ -117,7 +120,9 @@ impl Tree {
 
                 while let Ok(TraversalState::Ongoing(node)) = rx.recv() {
                     if let Some(ref mailbox) = progress_indicator_mailbox {
-                        let _ = mailbox.send(Message::Index);
+                        if mailbox.send(Message::Index).is_err() {
+                            return Err(Error::Terminated);
+                        }
                     }
 
                     if node.is_dir() {
@@ -147,7 +152,9 @@ impl Tree {
                 }
 
                 if let Some(ref mailbox) = progress_indicator_mailbox {
-                    let _ = mailbox.send(Message::DoneIndexing);
+                    if mailbox.send(Message::DoneIndexing).is_err() {
+                        return Err(Error::Terminated);
+                    }
                 }
 
                 let root_id = root_id.ok_or(Error::MissingRoot)?;
@@ -179,7 +186,7 @@ impl Tree {
 
             walker.visit(&mut visitor_builder);
 
-            tx.send(TraversalState::Done).unwrap();
+            let _ = tx.send(TraversalState::Done);
 
             res.join().unwrap()
         })
diff --git a/src/tree/visitor.rs b/src/tree/visitor.rs
index f0ff955..add11a9 100644
--- a/src/tree/visitor.rs
+++ b/src/tree/visitor.rs
@@ -44,7 +44,9 @@ impl ParallelVisitor for Branch<'_> {
 
         match Node::try_from((dir_entry, self.ctx)) {
             Ok(node) => {
-                self.tx.send(TraversalState::from(node)).unwrap();
+                if self.tx.send(TraversalState::from(node)).is_err() {
+                    return WalkState::Quit;
+                }
                 WalkState::Continue
             },
             _ => WalkState::Skip,
diff --git a/src/tty/mod.rs b/src/tty/mod.rs
index 502c3d8..b11da25 100644
--- a/src/tty/mod.rs
+++ b/src/tty/mod.rs
@@ -1,4 +1,5 @@
 #![allow(clippy::module_name_repetitions)]
+use crossterm::{cursor, ExecutableCommand};
 use std::io::{stdin, stdout, IsTerminal};
 
 #[cfg(windows)]
@@ -18,6 +19,11 @@ pub fn stdout_is_tty() -> bool {
     stdout().is_terminal()
 }
 
+/// Restore terminal settings.
+pub fn restore_tty() {
+    stdout().execute(cursor::Show).expect("Failed to restore cursor");
+}
+
 /// Attempts to get the current size of the tty's window. Returns `None` if stdout isn't tty or if
 /// failed to get width.
 pub fn get_window_width(stdout_is_tty: bool) -> Option<usize> {