Make formatting happen asynchronously.

4 files changed, 128 insertions, 40 deletions

diff --git a/helix-lsp/src/lib.rs b/helix-lsp/src/lib.rs
index e4ab153c..96a45bb9 100644
--- a/helix-lsp/src/lib.rs
+++ b/helix-lsp/src/lib.rs
@@ -182,6 +182,22 @@ pub mod util {
             }),
         )
     }
+
+    /// The result of asking the language server to format the document. This can be turned into a
+    /// `Transaction`, but the advantage of not doing that straight away is that this one is
+    /// `Send` and `Sync`.
+    #[derive(Clone, Debug)]
+    pub struct LspFormatting {
+        pub doc: Rope,
+        pub edits: Vec<lsp::TextEdit>,
+        pub offset_encoding: OffsetEncoding,
+    }
+
+    impl From<LspFormatting> for Transaction {
+        fn from(fmt: LspFormatting) -> Transaction {
+            generate_transaction_from_edits(&fmt.doc, fmt.edits, fmt.offset_encoding)
+        }
+    }
 }
 
 #[derive(Debug, PartialEq, Clone)]
diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index c3f0f392..710fc2a6 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -35,7 +35,7 @@ use crate::{
     ui::{self, Completion, Picker, Popup, Prompt, PromptEvent},
 };
 
-use crate::application::{LspCallbackWrapper, LspCallbacks};
+use crate::application::{LspCallback, LspCallbackWrapper, LspCallbacks};
 use futures_util::FutureExt;
 use std::{fmt, future::Future, path::Display, str::FromStr};
 
@@ -1145,11 +1145,12 @@ mod cmd {
     }
 
     fn write_impl<P: AsRef<Path>>(
-        view: &View,
-        doc: &mut Document,
+        cx: &mut compositor::Context,
         path: Option<P>,
     ) -> Result<tokio::task::JoinHandle<Result<(), anyhow::Error>>, anyhow::Error> {
         use anyhow::anyhow;
+        let callbacks = &mut cx.callbacks;
+        let (view, doc) = current!(cx.editor);
 
         if let Some(path) = path {
             if let Err(err) = doc.set_path(path.as_ref()) {
@@ -1163,15 +1164,21 @@ mod cmd {
             .language_config()
             .map(|config| config.auto_format)
             .unwrap_or_default();
-        if autofmt {
-            doc.format(view.id); // TODO: merge into save
-        }
-        Ok(tokio::spawn(doc.save()))
+        let fmt = if autofmt {
+            doc.format().map(|fmt| {
+                let shared = fmt.shared();
+                let callback = make_format_callback(doc.id(), doc.version(), true, shared.clone());
+                callbacks.push(callback);
+                shared
+            })
+        } else {
+            None
+        };
+        Ok(tokio::spawn(doc.format_and_save(fmt)))
     }
 
     fn write(cx: &mut compositor::Context, args: &[&str], event: PromptEvent) {
-        let (view, doc) = current!(cx.editor);
-        if let Err(e) = write_impl(view, doc, args.first()) {
+        if let Err(e) = write_impl(cx, args.first()) {
             cx.editor.set_error(e.to_string());
         };
     }
@@ -1181,9 +1188,12 @@ mod cmd {
     }
 
     fn format(cx: &mut compositor::Context, args: &[&str], event: PromptEvent) {
-        let (view, doc) = current!(cx.editor);
+        let (_, doc) = current!(cx.editor);
 
-        doc.format(view.id)
+        if let Some(format) = doc.format() {
+            let callback = make_format_callback(doc.id(), doc.version(), false, format);
+            cx.callbacks.push(callback);
+        }
     }
 
     fn set_indent_style(cx: &mut compositor::Context, args: &[&str], event: PromptEvent) {
@@ -1288,8 +1298,7 @@ mod cmd {
     }
 
     fn write_quit(cx: &mut compositor::Context, args: &[&str], event: PromptEvent) {
-        let (view, doc) = current!(cx.editor);
-        match write_impl(view, doc, args.first()) {
+        match write_impl(cx, args.first()) {
             Ok(handle) => {
                 if let Err(e) = helix_lsp::block_on(handle) {
                     cx.editor.set_error(e.to_string());
@@ -1305,7 +1314,7 @@ mod cmd {
 
     fn force_write_quit(cx: &mut compositor::Context, args: &[&str], event: PromptEvent) {
         let (view, doc) = current!(cx.editor);
-        match write_impl(view, doc, args.first()) {
+        match write_impl(cx, args.first()) {
             Ok(handle) => {
                 if let Err(e) = helix_lsp::block_on(handle) {
                     cx.editor.set_error(e.to_string());
@@ -1902,6 +1911,35 @@ fn append_to_line(cx: &mut Context) {
     doc.set_selection(view.id, selection);
 }
 
+// Creates an LspCallback that waits for formatting changes to be computed. When they're done,
+// it applies them, but only if the doc hasn't changed.
+fn make_format_callback(
+    doc_id: DocumentId,
+    doc_version: i32,
+    set_unmodified: bool,
+    format: impl Future<Output = helix_lsp::util::LspFormatting> + Send + 'static,
+) -> LspCallback {
+    Box::pin(async move {
+        let format = format.await;
+        let call: LspCallbackWrapper =
+            Box::new(move |editor: &mut Editor, compositor: &mut Compositor| {
+                let view_id = view!(editor).id;
+                if let Some(doc) = editor.document_mut(doc_id) {
+                    if doc.version() == doc_version {
+                        doc.apply(&Transaction::from(format), view_id);
+                        doc.append_changes_to_history(view_id);
+                        if set_unmodified {
+                            doc.reset_modified();
+                        }
+                    } else {
+                        log::info!("discarded formatting changes because the document changed");
+                    }
+                }
+            });
+        Ok(call)
+    })
+}
+
 enum Open {
     Below,
     Above,
diff --git a/helix-view/src/document.rs b/helix-view/src/document.rs
index dcacdb5e..80ef54d5 100644
--- a/helix-view/src/document.rs
+++ b/helix-view/src/document.rs
@@ -16,6 +16,7 @@ use helix_core::{
     ChangeSet, Diagnostic, LineEnding, Rope, RopeBuilder, Selection, State, Syntax, Transaction,
     DEFAULT_LINE_ENDING,
 };
+use helix_lsp::util::LspFormatting;
 
 use crate::{DocumentId, Theme, ViewId};
 
@@ -472,39 +473,55 @@ impl Document {
         Ok(doc)
     }
 
-    // TODO: remove view_id dependency here
-    pub fn format(&mut self, view_id: ViewId) {
-        if let Some(language_server) = self.language_server() {
-            // TODO: await, no blocking
-            let transaction = helix_lsp::block_on(language_server.text_document_formatting(
-                self.identifier(),
-                lsp::FormattingOptions::default(),
-                None,
-            ))
-            .map(|edits| {
-                helix_lsp::util::generate_transaction_from_edits(
-                    self.text(),
+    pub fn format(&self) -> Option<impl Future<Output = LspFormatting> + 'static> {
+        if let Some(language_server) = self.language_server.clone() {
+            let text = self.text.clone();
+            let id = self.identifier();
+            let fut = async move {
+                let edits = language_server
+                    .text_document_formatting(id, lsp::FormattingOptions::default(), None)
+                    .await
+                    .unwrap_or_else(|e| {
+                        log::warn!("LSP formatting failed: {}", e);
+                        Default::default()
+                    });
+                LspFormatting {
+                    doc: text,
                     edits,
-                    language_server.offset_encoding(),
-                )
-            });
-
-            if let Ok(transaction) = transaction {
-                self.apply(&transaction, view_id);
-                self.append_changes_to_history(view_id);
-            }
+                    offset_encoding: language_server.offset_encoding(),
+                }
+            };
+            Some(fut)
+        } else {
+            None
         }
     }
 
+    pub fn save(&mut self) -> impl Future<Output = Result<(), anyhow::Error>> {
+        self.save_impl::<futures_util::future::Ready<_>>(None)
+    }
+
+    pub fn format_and_save(
+        &mut self,
+        formatting: Option<impl Future<Output = LspFormatting>>,
+    ) -> impl Future<Output = anyhow::Result<()>> {
+        self.save_impl(formatting)
+    }
+
     // TODO: do we need some way of ensuring two save operations on the same doc can't run at once?
     // or is that handled by the OS/async layer
     /// The `Document`'s text is encoded according to its encoding and written to the file located
     /// at its `path()`.
-    pub fn save(&mut self) -> impl Future<Output = Result<(), anyhow::Error>> {
+    ///
+    /// If `formatting` is present, it supplies some changes that we apply to the text before saving.
+    fn save_impl<F: Future<Output = LspFormatting>>(
+        &mut self,
+        formatting: Option<F>,
+    ) -> impl Future<Output = Result<(), anyhow::Error>> {
         // we clone and move text + path into the future so that we asynchronously save the current
         // state without blocking any further edits.
 
-        let text = self.text().clone();
+        let mut text = self.text().clone();
         let path = self.path.clone().expect("Can't save with no path set!"); // TODO: handle no path
         let identifier = self.identifier();
 
@@ -512,10 +529,7 @@ impl Document {
 
         let language_server = self.language_server.clone();
 
-        // reset the modified flag
-        let history = self.history.take();
-        self.last_saved_revision = history.current_revision();
-        self.history.set(history);
+        self.reset_modified();
 
         let encoding = self.encoding;
 
@@ -531,6 +545,15 @@ impl Document {
                 }
             }
 
+            if let Some(fmt) = formatting {
+                let success = Transaction::from(fmt.await).changes().apply(&mut text);
+                if !success {
+                    // This shouldn't happen, because the transaction changes were generated
+                    // from the same text we're saving.
+                    log::error!("failed to apply format changes before saving");
+                }
+            }
+
             let mut file = File::create(path).await?;
             to_writer(&mut file, encoding, &text).await?;
 
@@ -877,6 +900,13 @@ impl Document {
         current_revision != self.last_saved_revision || !self.changes.is_empty()
     }
 
+    pub fn reset_modified(&mut self) {
+        let history = self.history.take();
+        let current_revision = history.current_revision();
+        self.history.set(history);
+        self.last_saved_revision = current_revision;
+    }
+
     pub fn mode(&self) -> Mode {
         self.mode
     }
diff --git a/helix-view/src/editor.rs b/helix-view/src/editor.rs
index 0a2dc54d..a16cc50f 100644
--- a/helix-view/src/editor.rs
+++ b/helix-view/src/editor.rs
@@ -270,6 +270,10 @@ impl Editor {
         self.documents.get(id)
     }
 
+    pub fn document_mut(&mut self, id: DocumentId) -> Option<&mut Document> {
+        self.documents.get_mut(id)
+    }
+
     pub fn documents(&self) -> impl Iterator<Item = &Document> {
         self.documents.iter().map(|(_id, doc)| doc)
     }