tinymist_query/analysis/
code_action.rs

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