feat: add range formatting support to the language server (#1984)

Co-authored-by: Myriad-Dreamin <camiyoru@gmail.com>

crates/tinymist/src/lsp.rs

@@ -1,6 +1,6 @@
 use std::sync::OnceLock;
 
-use lsp_types::request::WorkspaceConfiguration;
+use lsp_types::request::*;
 use lsp_types::*;
 use reflexo::ImmutPath;
 use request::{RegisterCapability, UnregisterCapability};
@@ -305,12 +305,20 @@ impl ServerState {
         }
 
         const FORMATTING_REGISTRATION_ID: &str = "formatting";
-        const DOCUMENT_FORMATTING_METHOD_ID: &str = "textDocument/formatting";
+        const RANGE_FORMATTING_REGISTRATION_ID: &str = "rangeFormatting";
 
         pub fn get_formatting_registration() -> Registration {
             Registration {
                 id: FORMATTING_REGISTRATION_ID.to_owned(),
-                method: DOCUMENT_FORMATTING_METHOD_ID.to_owned(),
+                method: Formatting::METHOD.to_owned(),
+                register_options: None,
+            }
+        }
+
+        pub fn get_range_formatting_registration() -> Registration {
+            Registration {
+                id: RANGE_FORMATTING_REGISTRATION_ID.to_owned(),
+                method: RangeFormatting::METHOD.to_owned(),
                 register_options: None,
             }
         }
@@ -318,22 +326,35 @@ impl ServerState {
         pub fn get_formatting_unregistration() -> Unregistration {
             Unregistration {
                 id: FORMATTING_REGISTRATION_ID.to_owned(),
-                method: DOCUMENT_FORMATTING_METHOD_ID.to_owned(),
+                method: Formatting::METHOD.to_owned(),
+            }
+        }
+
+        pub fn get_range_formatting_unregistration() -> Unregistration {
+            Unregistration {
+                id: RANGE_FORMATTING_REGISTRATION_ID.to_owned(),
+                method: RangeFormatting::METHOD.to_owned(),
             }
         }
 
         match (enable, self.formatter_registered) {
             (true, false) => {
                 log::trace!("registering formatter");
-                self.register_capability(vec![get_formatting_registration()])
-                    .inspect(|_| self.formatter_registered = enable)
-                    .context("could not register formatter")
+                self.register_capability(vec![
+                    get_formatting_registration(),
+                    get_range_formatting_registration(),
+                ])
+                .inspect(|_| self.formatter_registered = enable)
+                .context("could not register formatter")
             }
             (false, true) => {
                 log::trace!("unregistering formatter");
-                self.unregister_capability(vec![get_formatting_unregistration()])
-                    .inspect(|_| self.formatter_registered = enable)
-                    .context("could not unregister formatter")
+                self.unregister_capability(vec![
+                    get_formatting_unregistration(),
+                    get_range_formatting_unregistration(),
+                ])
+                .inspect(|_| self.formatter_registered = enable)
+                .context("could not unregister formatter")
             }
             _ => Ok(()),
         }

crates/tinymist/src/lsp/init.rs

@@ -102,6 +102,8 @@ impl Initializer for SuperInit {
         });
         let document_formatting_provider =
             (!const_config.doc_fmt_dynamic_registration).then_some(OneOf::Left(true));
+        let document_range_formatting_provider =
+            (!const_config.doc_fmt_dynamic_registration).then_some(OneOf::Left(true));
 
         let file_operations = const_config.notify_will_rename_files.then(|| {
             WorkspaceFileOperationsServerCapabilities {
@@ -194,6 +196,7 @@ impl Initializer for SuperInit {
                     file_operations,
                 }),
                 document_formatting_provider,
+                document_range_formatting_provider,
                 inlay_hint_provider: Some(OneOf::Left(true)),
                 code_action_provider: Some(CodeActionProviderCapability::Simple(true)),
                 code_lens_provider: Some(CodeLensOptions {

crates/tinymist/src/lsp/query.rs

@@ -99,6 +99,19 @@ impl ServerState {
         erased_response(self.formatter.run(source))
     }
 
+    pub(crate) fn range_formatting(
+        &mut self,
+        params: DocumentRangeFormattingParams,
+    ) -> ScheduleResult {
+        if matches!(self.config.formatter_mode, FormatterMode::Disable) {
+            return just_ok(serde_json::Value::Null);
+        }
+
+        let path: ImmutPath = as_path(params.text_document).as_path().into();
+        let source = self.query_source(path, Ok)?;
+        erased_response(self.formatter.run_on_range(source, params.range))
+    }
+
     pub(crate) fn inlay_hint(&mut self, params: InlayHintParams) -> ScheduleResult {
         let path = as_path(params.text_document);
         let range = params.range;

crates/tinymist/src/server.rs

@@ -255,6 +255,7 @@ impl ServerState {
             .with_request_::<DocumentSymbolRequest>(State::document_symbol)
             // Sync for low latency
             .with_request_::<Formatting>(State::formatting)
+            .with_request_::<RangeFormatting>(State::range_formatting)
             .with_request_::<SelectionRangeRequest>(State::selection_range)
             // latency insensitive
             .with_request_::<InlayHintRequest>(State::inlay_hint)

crates/tinymist/src/task/format.rs

@@ -2,9 +2,9 @@
 
 use std::iter::zip;
 
-use lsp_types::TextEdit;
+use lsp_types::{Range, TextEdit};
 use sync_ls::{just_future, SchedulableResponse};
-use tinymist_query::{to_lsp_range, PositionEncoding};
+use tinymist_query::{to_lsp_range, to_typst_range, PositionEncoding};
 use typst::syntax::Source;
 
 use super::SyncTaskFactory;
@@ -55,6 +55,31 @@ impl FormatTask {
             Ok(formatted.and_then(|formatted| calc_diff(src, formatted, c.position_encoding)))
         })
     }
+
+    pub fn run_on_range(
+        &self,
+        src: Source,
+        range: Range,
+    ) -> SchedulableResponse<Option<Vec<TextEdit>>> {
+        fn format_impl(src: Source, range: Range, c: &FormatUserConfig) -> Option<Vec<TextEdit>> {
+            let typst_range = to_typst_range(range, c.position_encoding, &src)?;
+
+            match &c.config {
+                FormatterConfig::Typstyle(config) => {
+                    let format_result = typstyle_core::Typstyle::new(config.as_ref().clone())
+                        .format_source_range(src.clone(), typst_range)
+                        .ok()?;
+                    let mut new_full_text = src.text().to_owned();
+                    new_full_text.replace_range(format_result.source_range, &format_result.content);
+                    calc_diff(src, new_full_text, c.position_encoding)
+                }
+                _ => None,
+            }
+        }
+
+        let c = self.factory.task();
+        just_future(async move { Ok(format_impl(src, range, &c)) })
+    }
 }
 
 /// A simple implementation of the diffing algorithm, borrowed from

tests/e2e/e2e/lsp.rs

@@ -22,7 +22,7 @@ fn test_lsp() {
         });
 
         let hash = replay_log(&root.join("neovim"));
-        insta::assert_snapshot!(hash, @"siphash128_13:9409ded3fa346d873f6ab39516a53958");
+        insta::assert_snapshot!(hash, @"siphash128_13:8da05d2505df482442dccd7c7b200542");
     }
 
     {