1 files changed, 41 insertions, 6 deletions
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;