Change focus to modified docs on quit (#3872)

* Change focus to modified docs on quit

When quitting with modified documents, automatically switch focus to
one of them.

* Update helix-term/src/commands/typed.rs

Co-authored-by: Poliorcetics <poliorcetics@users.noreply.github.com>

* Make it work with buffer-close-all and the like

* Cleanup

Use Cow instead of String, and rename DoesntExist -> DoesNotExist

Co-authored-by: Poliorcetics <poliorcetics@users.noreply.github.com>

3 files changed, 59 insertions, 25 deletions

diff --git a/helix-term/src/commands/typed.rs b/helix-term/src/commands/typed.rs
index c2c37291..96ff75c5 100644
--- a/helix-term/src/commands/typed.rs
+++ b/helix-term/src/commands/typed.rs
@@ -2,7 +2,7 @@ use std::ops::Deref;
 
 use super::*;
 
-use helix_view::editor::{Action, ConfigEvent};
+use helix_view::editor::{Action, CloseError, ConfigEvent};
 use ui::completers::{self, Completer};
 
 #[derive(Clone)]
@@ -71,8 +71,29 @@ fn buffer_close_by_ids_impl(
     doc_ids: &[DocumentId],
     force: bool,
 ) -> anyhow::Result<()> {
-    for &doc_id in doc_ids {
-        editor.close_document(doc_id, force)?;
+    let (modified_ids, modified_names): (Vec<_>, Vec<_>) = doc_ids
+        .iter()
+        .filter_map(|&doc_id| {
+            if let Err(CloseError::BufferModified(name)) = editor.close_document(doc_id, force) {
+                Some((doc_id, name))
+            } else {
+                None
+            }
+        })
+        .unzip();
+
+    if let Some(first) = modified_ids.first() {
+        let current = doc!(editor);
+        // If the current document is unmodified, and there are modified
+        // documents, switch focus to the first modified doc.
+        if !modified_ids.contains(&current.id()) {
+            editor.switch(*first, Action::Replace);
+        }
+        bail!(
+            "{} unsaved buffer(s) remaining: {:?}",
+            modified_names.len(),
+            modified_names
+        );
     }
 
     Ok(())
@@ -513,23 +534,26 @@ fn force_write_quit(
     force_quit(cx, &[], event)
 }
 
-/// Results an error if there are modified buffers remaining and sets editor error,
-/// otherwise returns `Ok(())`
+/// Results in an error if there are modified buffers remaining and sets editor
+/// error, otherwise returns `Ok(())`. If the current document is unmodified,
+/// and there are modified documents, switches focus to one of them.
 pub(super) fn buffers_remaining_impl(editor: &mut Editor) -> anyhow::Result<()> {
-    let modified: Vec<_> = editor
+    let (modified_ids, modified_names): (Vec<_>, Vec<_>) = editor
         .documents()
         .filter(|doc| doc.is_modified())
-        .map(|doc| {
-            doc.relative_path()
-                .map(|path| path.to_string_lossy().to_string())
-                .unwrap_or_else(|| SCRATCH_BUFFER_NAME.into())
-        })
-        .collect();
-    if !modified.is_empty() {
+        .map(|doc| (doc.id(), doc.display_name()))
+        .unzip();
+    if let Some(first) = modified_ids.first() {
+        let current = doc!(editor);
+        // If the current document is unmodified, and there are modified
+        // documents, switch focus to the first modified doc.
+        if !modified_ids.contains(&current.id()) {
+            editor.switch(*first, Action::Replace);
+        }
         bail!(
             "{} unsaved buffer(s) remaining: {:?}",
-            modified.len(),
-            modified
+            modified_names.len(),
+            modified_names
         );
     }
     Ok(())
diff --git a/helix-view/src/document.rs b/helix-view/src/document.rs
index 2ef99c6a..b6b2f664 100644
--- a/helix-view/src/document.rs
+++ b/helix-view/src/document.rs
@@ -5,6 +5,7 @@ use helix_core::auto_pairs::AutoPairs;
 use helix_core::Range;
 use serde::de::{self, Deserialize, Deserializer};
 use serde::Serialize;
+use std::borrow::Cow;
 use std::cell::Cell;
 use std::collections::HashMap;
 use std::fmt::Display;
@@ -1038,6 +1039,12 @@ impl Document {
             .map(helix_core::path::get_relative_path)
     }
 
+    pub fn display_name(&self) -> Cow<'static, str> {
+        self.relative_path()
+            .map(|path| path.to_string_lossy().to_string().into())
+            .unwrap_or_else(|| SCRATCH_BUFFER_NAME.into())
+    }
+
     // transact(Fn) ?
 
     // -- LSP methods
diff --git a/helix-view/src/editor.rs b/helix-view/src/editor.rs
index f21244a9..89759439 100644
--- a/helix-view/src/editor.rs
+++ b/helix-view/src/editor.rs
@@ -1,6 +1,6 @@
 use crate::{
     clipboard::{get_clipboard_provider, ClipboardProvider},
-    document::{Mode, SCRATCH_BUFFER_NAME},
+    document::Mode,
     graphics::{CursorKind, Rect},
     info::Info,
     input::KeyEvent,
@@ -28,7 +28,7 @@ use tokio::{
     time::{sleep, Duration, Instant, Sleep},
 };
 
-use anyhow::{bail, Error};
+use anyhow::Error;
 
 pub use helix_core::diagnostic::Severity;
 pub use helix_core::register::Registers;
@@ -711,6 +711,14 @@ pub enum Action {
     VerticalSplit,
 }
 
+/// Error thrown on failed document closed
+pub enum CloseError {
+    /// Document doesn't exist
+    DoesNotExist,
+    /// Buffer is modified
+    BufferModified(String),
+}
+
 impl Editor {
     pub fn new(
         mut area: Rect,
@@ -1070,19 +1078,14 @@ impl Editor {
         self._refresh();
     }
 
-    pub fn close_document(&mut self, doc_id: DocumentId, force: bool) -> anyhow::Result<()> {
+    pub fn close_document(&mut self, doc_id: DocumentId, force: bool) -> Result<(), CloseError> {
         let doc = match self.documents.get(&doc_id) {
             Some(doc) => doc,
-            None => bail!("document does not exist"),
+            None => return Err(CloseError::DoesNotExist),
         };
 
         if !force && doc.is_modified() {
-            bail!(
-                "buffer {:?} is modified",
-                doc.relative_path()
-                    .map(|path| path.to_string_lossy().to_string())
-                    .unwrap_or_else(|| SCRATCH_BUFFER_NAME.into())
-            );
+            return Err(CloseError::BufferModified(doc.display_name().into_owned()));
         }
 
         if let Some(language_server) = doc.language_server() {