tinymist_query/analysis/
code_action.rs

1//! Provides code actions for the document.
2
3use ecow::eco_format;
4use lsp_types::{ChangeAnnotation, CreateFile, CreateFileOptions};
5use regex::Regex;
6use tinymist_analysis::syntax::{
7    PreviousItem, SyntaxClass, adjust_expr, node_ancestors, previous_items,
8};
9use tinymist_std::path::{diff, unix_slash};
10use typst::syntax::Side;
11
12use super::get_link_exprs_in;
13use crate::analysis::LinkTarget;
14use crate::prelude::*;
15use crate::syntax::{InterpretMode, interpret_mode_at};
16
17/// Analyzes the document and provides code actions.
18pub struct CodeActionWorker<'a> {
19    /// The local analysis context to work with.
20    ctx: &'a mut LocalContext,
21    /// The source document to analyze.
22    source: Source,
23    /// The code actions to provide.
24    pub actions: Vec<CodeAction>,
25    /// The lazily calculated local URL to [`Self::source`].
26    local_url: OnceLock<Option<Url>>,
27}
28
29impl<'a> CodeActionWorker<'a> {
30    /// Creates a new color action worker.
31    pub fn new(ctx: &'a mut LocalContext, source: Source) -> Self {
32        Self {
33            ctx,
34            source,
35            actions: Vec::new(),
36            local_url: OnceLock::new(),
37        }
38    }
39
40    fn local_url(&self) -> Option<&Url> {
41        self.local_url
42            .get_or_init(|| self.ctx.uri_for_id(self.source.id()).ok())
43            .as_ref()
44    }
45
46    #[must_use]
47    fn local_edits(&self, edits: Vec<EcoSnippetTextEdit>) -> Option<EcoWorkspaceEdit> {
48        Some(EcoWorkspaceEdit {
49            changes: Some(HashMap::from_iter([(self.local_url()?.clone(), edits)])),
50            ..Default::default()
51        })
52    }
53
54    #[must_use]
55    fn local_edit(&self, edit: EcoSnippetTextEdit) -> Option<EcoWorkspaceEdit> {
56        self.local_edits(vec![edit])
57    }
58
59    pub(crate) fn autofix(
60        &mut self,
61        root: &LinkedNode<'_>,
62        range: &Range<usize>,
63        context: &lsp_types::CodeActionContext,
64    ) -> Option<()> {
65        if let Some(only) = &context.only
66            && !only.is_empty()
67            && !only
68                .iter()
69                .any(|kind| *kind == CodeActionKind::EMPTY || *kind == CodeActionKind::QUICKFIX)
70        {
71            return None;
72        }
73
74        for diag in &context.diagnostics {
75            if diag.source.as_ref().is_none_or(|t| t != "typst") {
76                continue;
77            }
78
79            match match_autofix_kind(diag.message.as_str()) {
80                Some(AutofixKind::UnknownVariable) => {
81                    self.autofix_unknown_variable(root, range);
82                }
83                Some(AutofixKind::FileNotFound) => {
84                    self.autofix_file_not_found(root, range);
85                }
86                _ => {}
87            }
88        }
89
90        Some(())
91    }
92
93    /// Automatically fixes unknown variable errors.
94    pub fn autofix_unknown_variable(
95        &mut self,
96        root: &LinkedNode,
97        range: &Range<usize>,
98    ) -> Option<()> {
99        let cursor = (range.start + 1).min(self.source.text().len());
100        let node = root.leaf_at_compat(cursor)?;
101        self.create_missing_variable(root, &node);
102        self.add_spaces_to_math_unknown_variable(&node);
103        Some(())
104    }
105
106    fn create_missing_variable(
107        &mut self,
108        root: &LinkedNode<'_>,
109        node: &LinkedNode<'_>,
110    ) -> Option<()> {
111        let ident = 'determine_ident: {
112            if let Some(ident) = node.cast::<ast::Ident>() {
113                break 'determine_ident ident.get().clone();
114            }
115            if let Some(ident) = node.cast::<ast::MathIdent>() {
116                break 'determine_ident ident.get().clone();
117            }
118
119            return None;
120        };
121
122        enum CreatePosition {
123            Before(usize),
124            After(usize),
125            Bad,
126        }
127
128        let previous_decl = previous_items(node.clone(), |item| {
129            match item {
130                PreviousItem::Parent(parent, ..) => match parent.kind() {
131                    SyntaxKind::LetBinding => {
132                        let mut create_before = parent.clone();
133                        while let Some(before) = create_before.prev_sibling() {
134                            if matches!(before.kind(), SyntaxKind::Hash) {
135                                create_before = before;
136                                continue;
137                            }
138
139                            break;
140                        }
141
142                        return Some(CreatePosition::Before(create_before.range().start));
143                    }
144                    SyntaxKind::CodeBlock | SyntaxKind::ContentBlock => {
145                        let child = parent.children().find(|child| {
146                            matches!(
147                                child.kind(),
148                                SyntaxKind::LeftBrace | SyntaxKind::LeftBracket
149                            )
150                        })?;
151
152                        return Some(CreatePosition::After(child.range().end));
153                    }
154                    SyntaxKind::ModuleImport | SyntaxKind::ModuleInclude => {
155                        return Some(CreatePosition::Bad);
156                    }
157                    _ => {}
158                },
159                PreviousItem::Sibling(node) => {
160                    if matches!(
161                        node.kind(),
162                        SyntaxKind::ModuleImport | SyntaxKind::ModuleInclude
163                    ) {
164                        // todo: hash
165                        return Some(CreatePosition::After(node.range().end));
166                    }
167                }
168            }
169
170            None
171        });
172
173        let (create_pos, side) = match previous_decl {
174            Some(CreatePosition::Before(pos)) => (pos, Side::Before),
175            Some(CreatePosition::After(pos)) => (pos, Side::After),
176            None => (0, Side::After),
177            Some(CreatePosition::Bad) => return None,
178        };
179
180        let pos_node = root.leaf_at(create_pos, side.clone());
181        let mode = match interpret_mode_at(pos_node.as_ref()) {
182            InterpretMode::Markup => "#",
183            _ => "",
184        };
185
186        let extend_assign = if self.ctx.analysis.extended_code_action {
187            " = ${1:none}$0"
188        } else {
189            ""
190        };
191        let new_text = if matches!(side, Side::Before) {
192            eco_format!("{mode}let {ident}{extend_assign}\n\n")
193        } else {
194            eco_format!("\n\n{mode}let {ident}{extend_assign}")
195        };
196
197        let range = self.ctx.to_lsp_range(create_pos..create_pos, &self.source);
198        let edit = self.local_edit(EcoSnippetTextEdit::new(range, new_text))?;
199        let action = CodeAction {
200            title: "Create missing variable".to_string(),
201            kind: Some(CodeActionKind::QUICKFIX),
202            edit: Some(edit),
203            ..CodeAction::default()
204        };
205        self.actions.push(action);
206        Some(())
207    }
208
209    /// Add spaces between letters in an unknown math identifier: `$xyz$` -> `$x
210    /// y z$`.
211    fn add_spaces_to_math_unknown_variable(&mut self, node: &LinkedNode<'_>) -> Option<()> {
212        let ident = node.cast::<ast::MathIdent>()?.get();
213
214        // Rewrite `a_ij` as `a_(i j)`, not `a_i j`.
215        // Likewise rewrite `ab/c` as `(a b)/c`, not `a b/c`.
216        let needs_parens = matches!(
217            node.parent_kind(),
218            Some(SyntaxKind::MathAttach | SyntaxKind::MathFrac)
219        );
220        let new_text = if needs_parens {
221            eco_format!("({})", ident.chars().join(" "))
222        } else {
223            ident.chars().join(" ").into()
224        };
225
226        let range = self.ctx.to_lsp_range(node.range(), &self.source);
227        let edit = self.local_edit(EcoSnippetTextEdit::new(range, new_text))?;
228        let action = CodeAction {
229            title: "Add spaces between letters".to_string(),
230            kind: Some(CodeActionKind::QUICKFIX),
231            edit: Some(edit),
232            ..CodeAction::default()
233        };
234        self.actions.push(action);
235        Some(())
236    }
237
238    /// Automatically fixes file not found errors.
239    pub fn autofix_file_not_found(
240        &mut self,
241        root: &LinkedNode,
242        range: &Range<usize>,
243    ) -> Option<()> {
244        let cursor = (range.start + 1).min(self.source.text().len());
245        let node = root.leaf_at_compat(cursor)?;
246
247        let importing = node.cast::<ast::Str>()?.get();
248        if importing.starts_with('@') {
249            // todo: create local package?
250            // if importing.starts_with("@local") { return None; }
251
252            // This is a package import, not a file import.
253            return None;
254        }
255
256        let file_id = node.span().id()?;
257        let root_path = self.ctx.path_for_id(file_id.join("/")).ok()?;
258        let path_in_workspace = file_id.vpath().join(importing.as_str());
259        let new_path = path_in_workspace.resolve(root_path.as_path())?;
260        let new_file_url = path_to_url(&new_path).ok()?;
261
262        let edit = self.create_file(new_file_url, false);
263
264        let file_to_create = unix_slash(path_in_workspace.as_rooted_path());
265        let action = CodeAction {
266            title: format!("Create missing file at `{file_to_create}`"),
267            kind: Some(CodeActionKind::QUICKFIX),
268            edit: Some(edit),
269            ..CodeAction::default()
270        };
271        self.actions.push(action);
272
273        Some(())
274    }
275
276    /// Starts to work.
277    pub fn scoped(&mut self, root: &LinkedNode, range: &Range<usize>) -> Option<()> {
278        let cursor = (range.start + 1).min(self.source.text().len());
279        let node = root.leaf_at_compat(cursor)?;
280        let mut node = &node;
281
282        let mut heading_resolved = false;
283        let mut equation_resolved = false;
284        let mut path_resolved = false;
285
286        self.wrap_actions(node, range);
287
288        loop {
289            match node.kind() {
290                // Only the deepest heading is considered
291                SyntaxKind::Heading if !heading_resolved => {
292                    heading_resolved = true;
293                    self.heading_actions(node);
294                }
295                // Only the deepest equation is considered
296                SyntaxKind::Equation if !equation_resolved => {
297                    equation_resolved = true;
298                    self.equation_actions(node);
299                }
300                SyntaxKind::Str if !path_resolved => {
301                    path_resolved = true;
302                    self.path_actions(node, cursor);
303                }
304                _ => {}
305            }
306
307            node = node.parent()?;
308        }
309    }
310
311    fn path_actions(&mut self, node: &LinkedNode, cursor: usize) -> Option<()> {
312        // We can only process the case where the import path is a string.
313        if let Some(SyntaxClass::IncludePath(path_node) | SyntaxClass::ImportPath(path_node)) =
314            classify_syntax(node.clone(), cursor)
315        {
316            let str_node = adjust_expr(path_node)?;
317            let str_ast = str_node.cast::<ast::Str>()?;
318            return self.path_rewrite(self.source.id(), &str_ast.get(), &str_node);
319        }
320
321        let link_parent = node_ancestors(node)
322            .find(|node| matches!(node.kind(), SyntaxKind::FuncCall))
323            .unwrap_or(node);
324
325        // Actually there should be only one link left
326        if let Some(link_info) = get_link_exprs_in(link_parent) {
327            let objects = link_info.objects.into_iter();
328            let object_under_node = objects.filter(|link| link.range.contains(&cursor));
329
330            let mut resolved = false;
331            for link in object_under_node {
332                if let LinkTarget::Path(id, path) = link.target {
333                    // todo: is there a link that is not a path string?
334                    resolved = self.path_rewrite(id, &path, node).is_some() || resolved;
335                }
336            }
337
338            return resolved.then_some(());
339        }
340
341        None
342    }
343
344    /// Rewrites absolute paths from/to relative paths.
345    fn path_rewrite(&mut self, id: TypstFileId, path: &str, node: &LinkedNode) -> Option<()> {
346        if !matches!(node.kind(), SyntaxKind::Str) {
347            log::warn!("bad path node kind on code action: {:?}", node.kind());
348            return None;
349        }
350
351        let path = Path::new(path);
352
353        if path.starts_with("/") {
354            // Convert absolute path to relative path
355            let cur_path = id.vpath().as_rooted_path().parent().unwrap();
356            let new_path = diff(path, cur_path)?;
357            let edit = self.edit_str(node, unix_slash(&new_path))?;
358            let action = CodeAction {
359                title: "Convert to relative path".to_string(),
360                kind: Some(CodeActionKind::REFACTOR_REWRITE),
361                edit: Some(edit),
362                ..CodeAction::default()
363            };
364            self.actions.push(action);
365        } else {
366            // Convert relative path to absolute path
367            let mut new_path = id.vpath().as_rooted_path().parent().unwrap().to_path_buf();
368            for i in path.components() {
369                match i {
370                    std::path::Component::ParentDir => {
371                        new_path.pop().then_some(())?;
372                    }
373                    std::path::Component::Normal(name) => {
374                        new_path.push(name);
375                    }
376                    _ => {}
377                }
378            }
379            let edit = self.edit_str(node, unix_slash(&new_path))?;
380            let action = CodeAction {
381                title: "Convert to absolute path".to_string(),
382                kind: Some(CodeActionKind::REFACTOR_REWRITE),
383                edit: Some(edit),
384                ..CodeAction::default()
385            };
386            self.actions.push(action);
387        }
388
389        Some(())
390    }
391
392    fn edit_str(&mut self, node: &LinkedNode, new_content: String) -> Option<EcoWorkspaceEdit> {
393        if !matches!(node.kind(), SyntaxKind::Str) {
394            log::warn!("edit_str only works on string AST nodes: {:?}", node.kind());
395            return None;
396        }
397
398        self.local_edit(EcoSnippetTextEdit::new_plain(
399            self.ctx.to_lsp_range(node.range(), &self.source),
400            // todo: this is merely occasionally correct, abusing string escape (`fmt::Debug`)
401            eco_format!("{new_content:?}"),
402        ))
403    }
404
405    fn wrap_actions(&mut self, node: &LinkedNode, range: &Range<usize>) -> Option<()> {
406        if range.is_empty() {
407            return None;
408        }
409
410        let start_mode = interpret_mode_at(Some(node));
411        if !matches!(start_mode, InterpretMode::Markup | InterpretMode::Math) {
412            return None;
413        }
414
415        let edit = self.local_edits(vec![
416            EcoSnippetTextEdit::new_plain(
417                self.ctx
418                    .to_lsp_range(range.start..range.start, &self.source),
419                EcoString::inline("#["),
420            ),
421            EcoSnippetTextEdit::new_plain(
422                self.ctx.to_lsp_range(range.end..range.end, &self.source),
423                EcoString::inline("]"),
424            ),
425        ])?;
426
427        let action = CodeAction {
428            title: "Wrap with content block".to_string(),
429            kind: Some(CodeActionKind::REFACTOR_REWRITE),
430            edit: Some(edit),
431            ..CodeAction::default()
432        };
433        self.actions.push(action);
434
435        Some(())
436    }
437
438    fn heading_actions(&mut self, node: &LinkedNode) -> Option<()> {
439        let heading = node.cast::<ast::Heading>()?;
440        let depth = heading.depth().get();
441
442        // Only the marker is replaced, for minimal text change
443        let marker = node
444            .children()
445            .find(|child| child.kind() == SyntaxKind::HeadingMarker)?;
446        let marker_range = marker.range();
447
448        if depth > 1 {
449            // Decrease depth of heading
450            let action = CodeAction {
451                title: "Decrease depth of heading".to_string(),
452                kind: Some(CodeActionKind::REFACTOR_REWRITE),
453                edit: Some(self.local_edit(EcoSnippetTextEdit::new_plain(
454                    self.ctx.to_lsp_range(marker_range.clone(), &self.source),
455                    EcoString::inline("=").repeat(depth - 1),
456                ))?),
457                ..CodeAction::default()
458            };
459            self.actions.push(action);
460        }
461
462        // Increase depth of heading
463        let action = CodeAction {
464            title: "Increase depth of heading".to_string(),
465            kind: Some(CodeActionKind::REFACTOR_REWRITE),
466            edit: Some(self.local_edit(EcoSnippetTextEdit::new_plain(
467                self.ctx.to_lsp_range(marker_range, &self.source),
468                EcoString::inline("=").repeat(depth + 1),
469            ))?),
470            ..CodeAction::default()
471        };
472        self.actions.push(action);
473
474        Some(())
475    }
476
477    fn equation_actions(&mut self, node: &LinkedNode) -> Option<()> {
478        let equation = node.cast::<ast::Equation>()?;
479        let body = equation.body();
480        let is_block = equation.block();
481
482        let body = node.find(body.span())?;
483        let body_range = body.range();
484        let node_end = node.range().end;
485
486        let mut chs = node.children();
487        let chs = chs.by_ref();
488        let is_dollar = |node: &LinkedNode| node.kind() == SyntaxKind::Dollar;
489        let first_dollar = chs.take(1).find(is_dollar)?;
490        let last_dollar = chs.rev().take(1).find(is_dollar)?;
491
492        // Erroneous equation is skipped.
493        // For example, some unclosed equation.
494        if first_dollar.offset() == last_dollar.offset() {
495            return None;
496        }
497
498        let front_range = self
499            .ctx
500            .to_lsp_range(first_dollar.range().end..body_range.start, &self.source);
501        let back_range = self
502            .ctx
503            .to_lsp_range(body_range.end..last_dollar.range().start, &self.source);
504
505        // Retrieve punctuation to move
506        let mark_after_equation = self
507            .source
508            .text()
509            .get(node_end..)
510            .and_then(|text| {
511                let mut ch = text.chars();
512                let nx = ch.next()?;
513                Some((nx, ch.next()))
514            })
515            .filter(|(ch, ch_next)| {
516                static IS_PUNCTUATION: LazyLock<Regex> =
517                    LazyLock::new(|| Regex::new(r"\p{Punctuation}").unwrap());
518                (ch.is_ascii_punctuation()
519                    && ch_next.is_none_or(|ch_next| !ch_next.is_ascii_punctuation()))
520                    || (!ch.is_ascii_punctuation() && IS_PUNCTUATION.is_match(&ch.to_string()))
521            });
522        let punc_modify = if let Some((nx, _)) = mark_after_equation {
523            let ch_range = self
524                .ctx
525                .to_lsp_range(node_end..node_end + nx.len_utf8(), &self.source);
526            let remove_edit = EcoSnippetTextEdit::new_plain(ch_range, EcoString::new());
527            Some((nx, remove_edit))
528        } else {
529            None
530        };
531
532        let rewrite_action = |title: &str, new_text: &str| {
533            let mut edits = vec![
534                EcoSnippetTextEdit::new_plain(front_range, new_text.into()),
535                EcoSnippetTextEdit::new_plain(
536                    back_range,
537                    if !new_text.is_empty() {
538                        if let Some((ch, _)) = &punc_modify {
539                            EcoString::from(*ch) + new_text
540                        } else {
541                            new_text.into()
542                        }
543                    } else {
544                        EcoString::new()
545                    },
546                ),
547            ];
548
549            if !new_text.is_empty()
550                && let Some((_, edit)) = &punc_modify
551            {
552                edits.push(edit.clone());
553            }
554
555            Some(CodeAction {
556                title: title.to_owned(),
557                kind: Some(CodeActionKind::REFACTOR_REWRITE),
558                edit: Some(self.local_edits(edits)?),
559                ..CodeAction::default()
560            })
561        };
562
563        // Prepare actions
564        let toggle_action = if is_block {
565            rewrite_action("Convert to inline equation", "")?
566        } else {
567            rewrite_action("Convert to block equation", " ")?
568        };
569        let block_action = rewrite_action("Convert to multiple-line block equation", "\n");
570
571        self.actions.push(toggle_action);
572        if let Some(a2) = block_action {
573            self.actions.push(a2);
574        }
575
576        Some(())
577    }
578
579    fn create_file(&self, uri: Url, needs_confirmation: bool) -> EcoWorkspaceEdit {
580        let change_id = "Typst Create Missing Files".to_string();
581
582        let create_op = EcoDocumentChangeOperation::Op(lsp_types::ResourceOp::Create(CreateFile {
583            uri,
584            options: Some(CreateFileOptions {
585                overwrite: Some(false),
586                ignore_if_exists: None,
587            }),
588            annotation_id: Some(change_id.clone()),
589        }));
590
591        let mut change_annotations = HashMap::new();
592        change_annotations.insert(
593            change_id.clone(),
594            ChangeAnnotation {
595                label: change_id,
596                needs_confirmation: Some(needs_confirmation),
597                description: Some("The file is missing but required by code".to_string()),
598            },
599        );
600
601        EcoWorkspaceEdit {
602            changes: None,
603            document_changes: Some(EcoDocumentChanges::Operations(vec![create_op])),
604            change_annotations: Some(change_annotations),
605        }
606    }
607}
608
609#[derive(Debug, Clone, Copy)]
610enum AutofixKind {
611    UnknownVariable,
612    FileNotFound,
613}
614
615fn match_autofix_kind(msg: &str) -> Option<AutofixKind> {
616    static PATTERNS: &[(&str, AutofixKind)] = &[
617        ("unknown variable", AutofixKind::UnknownVariable), // typst compiler error
618        ("file not found", AutofixKind::FileNotFound),
619    ];
620
621    for (pattern, kind) in PATTERNS {
622        if msg.starts_with(pattern) {
623            return Some(*kind);
624        }
625    }
626
627    None
628}