Allow initial scan to be interrupted properly…

…which requires peeling through all the layers, but it's worth it.

5 files changed, 67 insertions, 29 deletions

diff --git a/src/common.rs b/src/common.rs
index b405ec0..f88c491 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -189,3 +189,13 @@ pub struct WalkResult {
     /// The amount of io::errors we encountered. Can happen when fetching meta-data, or when reading the directory contents.
     pub num_errors: u64,
 }
+
+impl WalkResult {
+    pub fn to_exit_code(&self) -> i32 {
+        if self.num_errors > 0 {
+            1
+        } else {
+            0
+        }
+    }
+}
diff --git a/src/interactive/app/eventloop.rs b/src/interactive/app/eventloop.rs
index c7e0945..114bd1a 100644
--- a/src/interactive/app/eventloop.rs
+++ b/src/interactive/app/eventloop.rs
@@ -38,6 +38,11 @@ pub struct AppState {
     pub is_scanning: bool,
 }
 
+pub enum ProcessingResult {
+    Finished(WalkResult),
+    ExitRequested(WalkResult),
+}
+
 impl AppState {
     pub fn draw<B>(
         &mut self,
@@ -64,7 +69,7 @@ impl AppState {
         display: &mut DisplayOptions,
         terminal: &mut Terminal<B>,
         keys: impl Iterator<Item = Result<Key, io::Error>>,
-    ) -> Result<WalkResult, Error>
+    ) -> Result<ProcessingResult, Error>
     where
         B: Backend,
     {
@@ -79,9 +84,17 @@ impl AppState {
                 Char('\t') => {
                     self.cycle_focus(window);
                 }
-                Ctrl('c') => break,
+                Ctrl('c') => {
+                    return Ok(ProcessingResult::ExitRequested(WalkResult {
+                        num_errors: traversal.io_errors,
+                    }))
+                }
                 Char('q') | Esc => match self.focussed {
-                    Main => break,
+                    Main => {
+                        return Ok(ProcessingResult::ExitRequested(WalkResult {
+                            num_errors: traversal.io_errors,
+                        }))
+                    }
                     Mark => self.focussed = Main,
                     Help => {
                         self.focussed = Main;
@@ -119,9 +132,9 @@ impl AppState {
             };
             self.draw(window, traversal, display.clone(), terminal)?;
         }
-        Ok(WalkResult {
+        Ok(ProcessingResult::Finished(WalkResult {
             num_errors: traversal.io_errors,
-        })
+        }))
     }
 }
 
@@ -156,13 +169,15 @@ impl TerminalApp {
     where
         B: Backend,
     {
-        self.state.process_events(
+        match self.state.process_events(
             &mut self.window,
             &mut self.traversal,
             &mut self.display,
             terminal,
             keys,
-        )
+        )? {
+            ProcessingResult::Finished(res) | ProcessingResult::ExitRequested(res) => Ok(res),
+        }
     }
 
     pub fn initialize<B>(
@@ -170,14 +185,14 @@ impl TerminalApp {
         options: WalkOptions,
         input: Vec<PathBuf>,
         mode: Interaction,
-    ) -> Result<TerminalApp, Error>
+    ) -> Result<Option<TerminalApp>, Error>
     where
         B: Backend,
     {
         terminal.hide_cursor()?;
         terminal.clear()?;
         let mut display: DisplayOptions = options.clone().into();
-        display.byte_vis = ByteVisualization::Bar;
+        display.byte_vis = ByteVisualization::PercentageAndBar;
         let mut window = MainWindow::default();
         let (keys_tx, keys_rx) = crossbeam_channel::unbounded();
         match mode {
@@ -207,10 +222,13 @@ impl TerminalApp {
                 None => {
                     state = Some({
                         let sorting = Default::default();
+                        let entries =
+                            sorted_entries(&traversal.tree, traversal.root_index, sorting);
                         AppState {
                             root: traversal.root_index,
                             sorting,
-                            entries: sorted_entries(&traversal.tree, traversal.root_index, sorting),
+                            selected: entries.get(0).map(|b| b.index),
+                            entries,
                             is_scanning: true,
                             ..Default::default()
                         }
@@ -218,19 +236,25 @@ impl TerminalApp {
                     state.as_mut().expect("state to be present, we just set it")
                 }
             };
-            s.process_events(
+            let should_exit = match s.process_events(
                 &mut window,
                 traversal,
                 &mut display,
                 terminal,
                 fetch_buffered_key_events().into_iter(),
-            )?;
-            Ok(())
+            )? {
+                ProcessingResult::ExitRequested(_) => true,
+                ProcessingResult::Finished(_) => false,
+            };
+            Ok(should_exit)
         })?;
+        let traversal = match traversal {
+            Some(t) => t,
+            None => return Ok(None),
+        };
         drop(fetch_buffered_key_events); // shutdown input event handler early for good measure
 
-        display.byte_vis = ByteVisualization::PercentageAndBar;
-        Ok(TerminalApp {
+        Ok(Some(TerminalApp {
             state: {
                 let mut s = state.unwrap_or_else(|| {
                     let sorting = Default::default();
@@ -251,7 +275,7 @@ impl TerminalApp {
             display,
             traversal,
             window,
-        })
+        }))
     }
 }
 
diff --git a/src/interactive/app_test/utils.rs b/src/interactive/app_test/utils.rs
index 4aabf18..e46f9e7 100644
--- a/src/interactive/app_test/utils.rs
+++ b/src/interactive/app_test/utils.rs
@@ -179,7 +179,10 @@ pub fn initialized_app_and_terminal_with_closure<P: AsRef<Path>>(
         input,
         Interaction::None,
     )?;
-    Ok((terminal, app))
+    Ok((
+        terminal,
+        app.expect("app that didn't try to abort iteration"),
+    ))
 }
 
 pub fn new_test_terminal() -> std::io::Result<Terminal<TestBackend>> {
diff --git a/src/main.rs b/src/main.rs
index a4a2227..9c42272 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -43,18 +43,20 @@ fn run() -> Result<(), Error> {
                 let backend = TermionBackend::new(stdout);
                 Terminal::new(backend)?
             };
-            let mut app = TerminalApp::initialize(
+            let res = TerminalApp::initialize(
                 &mut terminal,
                 walk_options,
                 paths_from(input)?,
                 Interaction::Full,
-            )?;
-            app.process_events(&mut terminal, io::stdin().keys())?;
+            )?
+            .map(|mut app| app.process_events(&mut terminal, io::stdin().keys()));
+
             drop(terminal);
             io::stdout().flush().ok();
+
             // Exit 'quickly' to avoid having to wait for all memory to be freed by us.
             // Let the OS do it - we have nothing to lose, literally.
-            std::process::exit(0);
+            std::process::exit(res.transpose()?.map(|e| e.to_exit_code()).unwrap_or(0));
         }
         Some(Aggregate {
             input,
@@ -90,10 +92,7 @@ fn run() -> Result<(), Error> {
         }
     };
 
-    if res.num_errors > 0 {
-        process::exit(1);
-    }
-    Ok(())
+    process::exit(res.to_exit_code());
 }
 
 fn paths_from(paths: Vec<PathBuf>) -> Result<Vec<PathBuf>, io::Error> {
diff --git a/src/traverse.rs b/src/traverse.rs
index 8d4d060..1e2e1d3 100644
--- a/src/traverse.rs
+++ b/src/traverse.rs
@@ -37,8 +37,8 @@ impl Traversal {
     pub fn from_walk(
         mut walk_options: WalkOptions,
         input: Vec<PathBuf>,
-        mut update: impl FnMut(&mut Traversal) -> Result<(), Error>,
-    ) -> Result<Traversal, Error> {
+        mut update: impl FnMut(&mut Traversal) -> Result<bool, Error>,
+    ) -> Result<Option<Traversal>, Error> {
         fn set_size_or_panic(tree: &mut Tree, node_idx: TreeIndex, current_size_at_depth: u64) {
             tree.node_weight_mut(node_idx)
                 .expect("node for parent index we just retrieved")
@@ -178,7 +178,9 @@ impl Traversal {
                     last_seen_eid = eid;
                     last_checked = now;
 
-                    update(&mut t)?;
+                    if update(&mut t)? {
+                        return Ok(None);
+                    }
                 }
             }
         }
@@ -194,7 +196,7 @@ impl Traversal {
         set_size_or_panic(&mut t.tree, t.root_index, root_size);
         t.total_bytes = Some(root_size);
 
-        Ok(t)
+        Ok(Some(t))
     }
 
     fn recompute_root_size(&self) -> u64 {