feat: provide code action to rewrite headings (#240)

* feat: provide code action to rewrite headings

* dev: update snapshot

crates/tinymist-query/src/code_action.rs

@@ -0,0 +1,130 @@
+use lsp_types::TextEdit;
+use once_cell::sync::OnceCell;
+
+use crate::{prelude::*, SemanticRequest};
+
+/// The [`textDocument/codeLens`] request is sent from the client to the server
+/// to compute code lenses for a given text document.
+///
+/// [`textDocument/codeLens`]: https://microsoft.github.io/language-server-protocol/specification#textDocument_codeLens
+#[derive(Debug, Clone)]
+pub struct CodeActionRequest {
+    /// The path of the document to request for.
+    pub path: PathBuf,
+    /// The range of the document to get code actions for.
+    pub range: LspRange,
+}
+
+impl SemanticRequest for CodeActionRequest {
+    type Response = Vec<CodeActionOrCommand>;
+
+    fn request(self, ctx: &mut AnalysisContext) -> Option<Self::Response> {
+        let source = ctx.source_by_path(&self.path).ok()?;
+        let range = ctx.to_typst_range(self.range, &source)?;
+        let cursor = (range.start + 1).min(source.text().len());
+        // todo: don't ignore the range end
+
+        let root = LinkedNode::new(source.root());
+        let mut worker = CodeActionWorker::new(ctx, source.clone());
+        worker.work(root, cursor);
+
+        let res = worker.actions;
+        (!res.is_empty()).then_some(res)
+    }
+}
+
+struct CodeActionWorker<'a, 'w> {
+    ctx: &'a mut AnalysisContext<'w>,
+    actions: Vec<CodeActionOrCommand>,
+    local_url: OnceCell<Option<Url>>,
+    current: Source,
+}
+
+impl<'a, 'w> CodeActionWorker<'a, 'w> {
+    fn new(ctx: &'a mut AnalysisContext<'w>, current: Source) -> Self {
+        Self {
+            ctx,
+            actions: Vec::new(),
+            local_url: OnceCell::new(),
+            current,
+        }
+    }
+
+    fn local_url(&self) -> Option<&Url> {
+        self.local_url
+            .get_or_init(|| {
+                let id = self.current.id();
+                let path = self.ctx.path_for_id(id).ok()?;
+                path_to_url(path.as_path()).ok()
+            })
+            .as_ref()
+    }
+
+    fn local_edit(&self, edit: TextEdit) -> Option<WorkspaceEdit> {
+        Some(WorkspaceEdit {
+            changes: Some(HashMap::from_iter([(
+                self.local_url()?.clone(),
+                vec![edit],
+            )])),
+            ..Default::default()
+        })
+    }
+
+    fn heading_actions(&mut self, leaf: &LinkedNode) -> Option<()> {
+        let h = leaf.cast::<ast::Heading>()?;
+        let depth = h.depth().get();
+
+        // Only the marker is replaced, for minimal text change
+        let marker = leaf
+            .children()
+            .find(|e| e.kind() == SyntaxKind::HeadingMarker)?;
+        let marker_range = marker.range();
+
+        if depth > 1 {
+            // decrease depth of heading
+            let action = CodeActionOrCommand::CodeAction(CodeAction {
+                title: "Decrease depth of heading".to_string(),
+                kind: Some(CodeActionKind::REFACTOR_REWRITE),
+                edit: Some(self.local_edit(TextEdit {
+                    range: self.ctx.to_lsp_range(marker_range.clone(), &self.current),
+                    new_text: "=".repeat(depth - 1),
+                })?),
+                ..CodeAction::default()
+            });
+
+            self.actions.push(action);
+        }
+
+        // increase depth of heading
+        let action = CodeActionOrCommand::CodeAction(CodeAction {
+            title: "Increase depth of heading".to_string(),
+            kind: Some(CodeActionKind::REFACTOR_REWRITE),
+            edit: Some(self.local_edit(TextEdit {
+                range: self.ctx.to_lsp_range(marker_range, &self.current),
+                new_text: "=".repeat(depth + 1),
+            })?),
+            ..CodeAction::default()
+        });
+
+        self.actions.push(action);
+
+        Some(())
+    }
+
+    fn work(&mut self, root: LinkedNode, cursor: usize) -> Option<()> {
+        let mut node = root.leaf_at(cursor)?;
+        let mut heading_resolved = false;
+        loop {
+            match node.kind() {
+                // Only the deepest heading is considered
+                SyntaxKind::Heading if !heading_resolved => {
+                    heading_resolved = true;
+                    self.heading_actions(&node);
+                }
+                _ => {}
+            }
+
+            node = node.parent()?.clone();
+        }
+    }
+}

crates/tinymist-query/src/lib.rs

@@ -20,6 +20,8 @@ pub use analysis::AnalysisContext;
 use typst::{model::Document as TypstDocument, syntax::Source};
 
 pub use diagnostics::*;
+pub(crate) mod code_action;
+pub use code_action::*;
 pub(crate) mod code_context;
 pub use code_context::*;
 pub(crate) mod code_lens;
@@ -205,6 +207,7 @@ mod polymorphic {
         InlayHint(InlayHintRequest),
         DocumentColor(DocumentColorRequest),
         ColorPresentation(ColorPresentationRequest),
+        CodeAction(CodeActionRequest),
         CodeLens(CodeLensRequest),
         Completion(CompletionRequest),
         SignatureHelp(SignatureHelpRequest),
@@ -236,6 +239,7 @@ mod polymorphic {
                 CompilerQueryRequest::InlayHint(..) => Unique,
                 CompilerQueryRequest::DocumentColor(..) => PinnedFirst,
                 CompilerQueryRequest::ColorPresentation(..) => ContextFreeUnique,
+                CompilerQueryRequest::CodeAction(..) => Unique,
                 CompilerQueryRequest::CodeLens(..) => Unique,
                 CompilerQueryRequest::Completion(..) => Mergeable,
                 CompilerQueryRequest::SignatureHelp(..) => PinnedFirst,
@@ -266,6 +270,7 @@ mod polymorphic {
                 CompilerQueryRequest::InlayHint(req) => &req.path,
                 CompilerQueryRequest::DocumentColor(req) => &req.path,
                 CompilerQueryRequest::ColorPresentation(req) => &req.path,
+                CompilerQueryRequest::CodeAction(req) => &req.path,
                 CompilerQueryRequest::CodeLens(req) => &req.path,
                 CompilerQueryRequest::Completion(req) => &req.path,
                 CompilerQueryRequest::SignatureHelp(req) => &req.path,
@@ -297,6 +302,7 @@ mod polymorphic {
         InlayHint(Option<Vec<InlayHint>>),
         DocumentColor(Option<Vec<ColorInformation>>),
         ColorPresentation(Option<Vec<ColorPresentation>>),
+        CodeAction(Option<Vec<CodeActionOrCommand>>),
         CodeLens(Option<Vec<CodeLens>>),
         Completion(Option<CompletionResponse>),
         SignatureHelp(Option<SignatureHelp>),

crates/tinymist-query/src/prelude.rs

@@ -9,13 +9,13 @@ pub use ecow::EcoVec;
 pub use itertools::{Format, Itertools};
 pub use log::{error, trace};
 pub use lsp_types::{
-    request::GotoDeclarationResponse, CodeLens, ColorInformation, ColorPresentation,
-    CompletionResponse, DiagnosticRelatedInformation, DocumentSymbol, DocumentSymbolResponse,
-    Documentation, FoldingRange, GotoDefinitionResponse, Hover, InlayHint, LanguageString,
-    Location as LspLocation, LocationLink, MarkedString, MarkupContent, MarkupKind,
-    Position as LspPosition, PrepareRenameResponse, SelectionRange, SemanticTokens,
-    SemanticTokensDelta, SemanticTokensFullDeltaResult, SemanticTokensResult, SignatureHelp,
-    SignatureInformation, SymbolInformation, Url, WorkspaceEdit,
+    request::GotoDeclarationResponse, CodeAction, CodeActionKind, CodeActionOrCommand, CodeLens,
+    ColorInformation, ColorPresentation, CompletionResponse, DiagnosticRelatedInformation,
+    DocumentSymbol, DocumentSymbolResponse, Documentation, FoldingRange, GotoDefinitionResponse,
+    Hover, InlayHint, LanguageString, Location as LspLocation, LocationLink, MarkedString,
+    MarkupContent, MarkupKind, Position as LspPosition, PrepareRenameResponse, SelectionRange,
+    SemanticTokens, SemanticTokensDelta, SemanticTokensFullDeltaResult, SemanticTokensResult,
+    SignatureHelp, SignatureInformation, SymbolInformation, Url, WorkspaceEdit,
 };
 pub use reflexo::vector::ir::DefId;
 pub use serde_json::Value as JsonValue;

crates/tinymist/src/server/lsp.rs

@@ -307,6 +307,7 @@ impl TypstLanguageServer {
             request_fn!(DocumentColor, Self::document_color),
             request_fn!(ColorPresentationRequest, Self::color_presentation),
             request_fn!(HoverRequest, Self::hover),
+            request_fn!(CodeActionRequest, Self::code_action),
             request_fn!(CodeLensRequest, Self::code_lens),
             request_fn!(FoldingRangeRequest, Self::folding_range),
             request_fn!(SignatureHelpRequest, Self::signature_help),
@@ -1195,6 +1196,15 @@ impl TypstLanguageServer {
         run_query!(self.ColorPresentation(path, color, range))
     }
 
+    fn code_action(
+        &mut self,
+        params: CodeActionParams,
+    ) -> LspResult<Option<Vec<CodeActionOrCommand>>> {
+        let path = as_path(params.text_document);
+        let range = params.range;
+        run_query!(self.CodeAction(path, range))
+    }
+
     fn code_lens(&mut self, params: CodeLensParams) -> LspResult<Option<Vec<CodeLens>>> {
         let path = as_path(params.text_document);
         run_query!(self.CodeLens(path))

crates/tinymist/src/server/lsp_init.rs

@@ -479,6 +479,7 @@ impl Init {
                 }),
                 document_formatting_provider,
                 inlay_hint_provider: Some(OneOf::Left(true)),
+                code_action_provider: Some(CodeActionProviderCapability::Simple(true)),
                 code_lens_provider: Some(CodeLensOptions {
                     resolve_provider: Some(false),
                 }),

crates/tinymist/src/state.rs

@@ -263,6 +263,7 @@ impl TypstLanguageServer {
             References(req) => query_world!(client, References, req),
             InlayHint(req) => query_world!(client, InlayHint, req),
             DocumentColor(req) => query_world!(client, DocumentColor, req),
+            CodeAction(req) => query_world!(client, CodeAction, req),
             CodeLens(req) => query_world!(client, CodeLens, req),
             Completion(req) => query_state!(client, Completion, req),
             SignatureHelp(req) => query_world!(client, SignatureHelp, req),

tests/e2e/main.rs

@@ -374,7 +374,7 @@ fn e2e() {
         });
 
         let hash = replay_log(&tinymist_binary, &root.join("neovim"));
-        insta::assert_snapshot!(hash, @"siphash128_13:e271d06cb84deee97e86b0b2763d741a");
+        insta::assert_snapshot!(hash, @"siphash128_13:8a283a742935063715fbc8af6bff2b88");
     }
 
     {
@@ -385,7 +385,7 @@ fn e2e() {
         });
 
         let hash = replay_log(&tinymist_binary, &root.join("vscode"));
-        insta::assert_snapshot!(hash, @"siphash128_13:86c5b1c533652019b3d99f7079e19270");
+        insta::assert_snapshot!(hash, @"siphash128_13:5f118cda0949759f157a2c22b9433775");
     }
 }