fix write-quit with auto format

write-quit will now save all files successfully even when there is auto
formatting

6 files changed, 85 insertions, 36 deletions

diff --git a/helix-term/src/application.rs b/helix-term/src/application.rs
index e84739cd..84eba22a 100644
--- a/helix-term/src/application.rs
+++ b/helix-term/src/application.rs
@@ -971,7 +971,11 @@ impl Application {
         //        want to try to run as much cleanup as we can, regardless of
         //        errors along the way
 
-        let mut result = match self.jobs.finish().await {
+        let mut result = match self
+            .jobs
+            .finish(Some(&mut self.editor), Some(&mut self.compositor))
+            .await
+        {
             Ok(_) => Ok(()),
             Err(err) => {
                 log::error!("Error executing job: {}", err);
diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index e76e0280..afd94564 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -47,6 +47,7 @@ use movement::Movement;
 use crate::{
     args,
     compositor::{self, Component, Compositor},
+    job::Callback,
     keymap::ReverseKeymap,
     ui::{self, overlay::overlayed, FilePicker, Picker, Popup, Prompt, PromptEvent},
 };
@@ -107,10 +108,11 @@ impl<'a> Context<'a> {
         let callback = Box::pin(async move {
             let json = call.await?;
             let response = serde_json::from_value(json)?;
-            let call: job::Callback =
-                Box::new(move |editor: &mut Editor, compositor: &mut Compositor| {
+            let call: job::Callback = Callback::EditorCompositor(Box::new(
+                move |editor: &mut Editor, compositor: &mut Compositor| {
                     callback(editor, compositor, response)
-                });
+                },
+            ));
             Ok(call)
         });
         self.jobs.callback(callback);
@@ -1925,8 +1927,8 @@ fn global_search(cx: &mut Context) {
     let show_picker = async move {
         let all_matches: Vec<FileResult> =
             UnboundedReceiverStream::new(all_matches_rx).collect().await;
-        let call: job::Callback =
-            Box::new(move |editor: &mut Editor, compositor: &mut Compositor| {
+        let call: job::Callback = Callback::EditorCompositor(Box::new(
+            move |editor: &mut Editor, compositor: &mut Compositor| {
                 if all_matches.is_empty() {
                     editor.set_status("No matches found");
                     return;
@@ -1962,7 +1964,8 @@ fn global_search(cx: &mut Context) {
                     },
                 );
                 compositor.push(Box::new(overlayed(picker)));
-            });
+            },
+        ));
         Ok(call)
     };
     cx.jobs.callback(show_picker);
@@ -2516,7 +2519,7 @@ async fn make_format_callback(
     write: Option<(Option<PathBuf>, bool)>,
 ) -> anyhow::Result<job::Callback> {
     let format = format.await?;
-    let call: job::Callback = Box::new(move |editor, _compositor| {
+    let call: job::Callback = Callback::EditorCompositor(Box::new(move |editor, _compositor| {
         if !editor.documents.contains_key(&doc_id) {
             return;
         }
@@ -2546,7 +2549,7 @@ async fn make_format_callback(
         } else {
             log::info!("discarded formatting changes because the document changed");
         }
-    });
+    }));
 
     Ok(call)
 }
diff --git a/helix-term/src/commands/dap.rs b/helix-term/src/commands/dap.rs
index 12a3fbc7..9c067eba 100644
--- a/helix-term/src/commands/dap.rs
+++ b/helix-term/src/commands/dap.rs
@@ -118,11 +118,14 @@ fn dap_callback<T, F>(
     let callback = Box::pin(async move {
         let json = call.await?;
         let response = serde_json::from_value(json)?;
-        let call: Callback = Box::new(move |editor: &mut Editor, compositor: &mut Compositor| {
-            callback(editor, compositor, response)
-        });
+        let call: Callback = Callback::EditorCompositor(Box::new(
+            move |editor: &mut Editor, compositor: &mut Compositor| {
+                callback(editor, compositor, response)
+            },
+        ));
         Ok(call)
     });
+
     jobs.callback(callback);
 }
 
@@ -274,10 +277,10 @@ pub fn dap_launch(cx: &mut Context) {
             let completions = template.completion.clone();
             let name = template.name.clone();
             let callback = Box::pin(async move {
-                let call: Callback = Box::new(move |_editor, compositor| {
+                let call: Callback = Callback::Compositor(Box::new(move |compositor| {
                     let prompt = debug_parameter_prompt(completions, name, Vec::new());
                     compositor.push(Box::new(prompt));
-                });
+                }));
                 Ok(call)
             });
             cx.jobs.callback(callback);
@@ -332,10 +335,10 @@ fn debug_parameter_prompt(
                 let config_name = config_name.clone();
                 let params = params.clone();
                 let callback = Box::pin(async move {
-                    let call: Callback = Box::new(move |_editor, compositor| {
+                    let call: Callback = Callback::Compositor(Box::new(move |compositor| {
                         let prompt = debug_parameter_prompt(completions, config_name, params);
                         compositor.push(Box::new(prompt));
-                    });
+                    }));
                     Ok(call)
                 });
                 cx.jobs.callback(callback);
@@ -582,7 +585,7 @@ pub fn dap_edit_condition(cx: &mut Context) {
             None => return,
         };
         let callback = Box::pin(async move {
-            let call: Callback = Box::new(move |editor, compositor| {
+            let call: Callback = Callback::EditorCompositor(Box::new(move |editor, compositor| {
                 let mut prompt = Prompt::new(
                     "condition:".into(),
                     None,
@@ -610,7 +613,7 @@ pub fn dap_edit_condition(cx: &mut Context) {
                     prompt.insert_str(&condition, editor)
                 }
                 compositor.push(Box::new(prompt));
-            });
+            }));
             Ok(call)
         });
         cx.jobs.callback(callback);
@@ -624,7 +627,7 @@ pub fn dap_edit_log(cx: &mut Context) {
             None => return,
         };
         let callback = Box::pin(async move {
-            let call: Callback = Box::new(move |editor, compositor| {
+            let call: Callback = Callback::EditorCompositor(Box::new(move |editor, compositor| {
                 let mut prompt = Prompt::new(
                     "log-message:".into(),
                     None,
@@ -651,7 +654,7 @@ pub fn dap_edit_log(cx: &mut Context) {
                     prompt.insert_str(&log_message, editor);
                 }
                 compositor.push(Box::new(prompt));
-            });
+            }));
             Ok(call)
         });
         cx.jobs.callback(callback);
diff --git a/helix-term/src/commands/typed.rs b/helix-term/src/commands/typed.rs
index 955b3b58..a687b854 100644
--- a/helix-term/src/commands/typed.rs
+++ b/helix-term/src/commands/typed.rs
@@ -519,9 +519,10 @@ fn write_quit(
 
     write_impl(cx, args.first(), false)?;
 
+    tokio::task::block_in_place(|| helix_lsp::block_on(cx.jobs.finish(Some(cx.editor), None)))?;
+
     let doc = doc_mut!(cx.editor);
 
-    tokio::task::block_in_place(|| helix_lsp::block_on(cx.jobs.finish()))?;
     tokio::task::block_in_place(|| helix_lsp::block_on(doc.try_flush_saves()))
         .map(|result| result.map(|_| ()))
         .unwrap_or(Ok(()))?;
@@ -1491,12 +1492,13 @@ fn tree_sitter_subtree(
             let contents = format!("```tsq\n{}\n```", selected_node.to_sexp());
 
             let callback = async move {
-                let call: job::Callback =
-                    Box::new(move |editor: &mut Editor, compositor: &mut Compositor| {
+                let call: job::Callback = Callback::EditorCompositor(Box::new(
+                    move |editor: &mut Editor, compositor: &mut Compositor| {
                         let contents = ui::Markdown::new(contents, editor.syn_loader.clone());
                         let popup = Popup::new("hover", contents).auto_close(true);
                         compositor.replace_or_push("hover", popup);
-                    });
+                    },
+                ));
                 Ok(call)
             };
 
@@ -1604,8 +1606,8 @@ fn run_shell_command(
 
     if !output.is_empty() {
         let callback = async move {
-            let call: job::Callback =
-                Box::new(move |editor: &mut Editor, compositor: &mut Compositor| {
+            let call: job::Callback = Callback::EditorCompositor(Box::new(
+                move |editor: &mut Editor, compositor: &mut Compositor| {
                     let contents = ui::Markdown::new(
                         format!("```sh\n{}\n```", output),
                         editor.syn_loader.clone(),
@@ -1614,7 +1616,8 @@ fn run_shell_command(
                         helix_core::Position::new(editor.cursor().0.unwrap_or_default().row, 2),
                     ));
                     compositor.replace_or_push("shell", popup);
-                });
+                },
+            ));
             Ok(call)
         };
 
diff --git a/helix-term/src/job.rs b/helix-term/src/job.rs
index e5147992..a997653d 100644
--- a/helix-term/src/job.rs
+++ b/helix-term/src/job.rs
@@ -5,7 +5,12 @@ use crate::compositor::Compositor;
 use futures_util::future::{BoxFuture, Future, FutureExt};
 use futures_util::stream::{FuturesUnordered, StreamExt};
 
-pub type Callback = Box<dyn FnOnce(&mut Editor, &mut Compositor) + Send>;
+pub enum Callback {
+    EditorCompositor(Box<dyn FnOnce(&mut Editor, &mut Compositor) + Send>),
+    Editor(Box<dyn FnOnce(&mut Editor) + Send>),
+    Compositor(Box<dyn FnOnce(&mut Compositor) + Send>),
+}
+
 pub type JobFuture = BoxFuture<'static, anyhow::Result<Option<Callback>>>;
 
 pub struct Job {
@@ -68,9 +73,11 @@ impl Jobs {
     ) {
         match call {
             Ok(None) => {}
-            Ok(Some(call)) => {
-                call(editor, compositor);
-            }
+            Ok(Some(call)) => match call {
+                Callback::EditorCompositor(call) => call(editor, compositor),
+                Callback::Editor(call) => call(editor),
+                Callback::Compositor(call) => call(compositor),
+            },
             Err(e) => {
                 editor.set_error(format!("Async job failed: {}", e));
             }
@@ -93,13 +100,41 @@ impl Jobs {
     }
 
     /// Blocks until all the jobs that need to be waited on are done.
-    pub async fn finish(&mut self) -> anyhow::Result<()> {
+    pub async fn finish(
+        &mut self,
+        mut editor: Option<&mut Editor>,
+        mut compositor: Option<&mut Compositor>,
+    ) -> anyhow::Result<()> {
         log::debug!("waiting on jobs...");
         let mut wait_futures = std::mem::take(&mut self.wait_futures);
         while let (Some(job), tail) = wait_futures.into_future().await {
             match job {
-                Ok(_) => {
+                Ok(callback) => {
                     wait_futures = tail;
+
+                    if let Some(callback) = callback {
+                        // clippy doesn't realize this is an error without the derefs
+                        #[allow(clippy::needless_option_as_deref)]
+                        match callback {
+                            Callback::EditorCompositor(call)
+                                if editor.is_some() && compositor.is_some() =>
+                            {
+                                call(
+                                    editor.as_deref_mut().unwrap(),
+                                    compositor.as_deref_mut().unwrap(),
+                                )
+                            }
+                            Callback::Editor(call) if editor.is_some() => {
+                                call(editor.as_deref_mut().unwrap())
+                            }
+                            Callback::Compositor(call) if compositor.is_some() => {
+                                call(compositor.as_deref_mut().unwrap())
+                            }
+
+                            // skip callbacks for which we don't have the necessary references
+                            _ => (),
+                        }
+                    }
                 }
                 Err(e) => {
                     self.wait_futures = tail;
diff --git a/helix-term/src/ui/editor.rs b/helix-term/src/ui/editor.rs
index 3cd2130a..abed7f9b 100644
--- a/helix-term/src/ui/editor.rs
+++ b/helix-term/src/ui/editor.rs
@@ -1,7 +1,8 @@
 use crate::{
     commands,
     compositor::{Component, Context, Event, EventResult},
-    job, key,
+    job::{self, Callback},
+    key,
     keymap::{KeymapResult, Keymaps},
     ui::{Completion, ProgressSpinners},
 };
@@ -944,9 +945,9 @@ impl EditorView {
 
                     // TODO: Use an on_mode_change hook to remove signature help
                     cxt.jobs.callback(async {
-                        let call: job::Callback = Box::new(|_editor, compositor| {
+                        let call: job::Callback = Callback::Compositor(Box::new(|compositor| {
                             compositor.remove(SignatureHelp::ID);
-                        });
+                        }));
                         Ok(call)
                     });
                 }