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 y z$`.
210    fn add_spaces_to_math_unknown_variable(&mut self, node: &LinkedNode<'_>) -> Option<()> {
211        let ident = node.cast::<ast::MathIdent>()?.get();
212
213        // Rewrite `a_ij` as `a_(i j)`, not `a_i j`.
214        // Likewise rewrite `ab/c` as `(a b)/c`, not `a b/c`.
215        let needs_parens = matches!(
216            node.parent_kind(),
217            Some(SyntaxKind::MathAttach | SyntaxKind::MathFrac)
218        );
219        let new_text = if needs_parens {
220            eco_format!("({})", ident.chars().join(" "))
221        } else {
222            ident.chars().join(" ").into()
223        };
224
225        let range = self.ctx.to_lsp_range(node.range(), &self.source);
226        let edit = self.local_edit(EcoSnippetTextEdit::new(range, new_text))?;
227        let action = CodeAction {
228            title: "Add spaces between letters".to_string(),
229            kind: Some(CodeActionKind::QUICKFIX),
230            edit: Some(edit),
231            ..CodeAction::default()
232        };
233        self.actions.push(action);
234        Some(())
235    }
236
237    /// Automatically fixes file not found errors.
238    pub fn autofix_file_not_found(
239        &mut self,
240        root: &LinkedNode,
241        range: &Range<usize>,
242    ) -> Option<()> {
243        let cursor = (range.start + 1).min(self.source.text().len());
244        let node = root.leaf_at_compat(cursor)?;
245
246        let importing = node.cast::<ast::Str>()?.get();
247        if importing.starts_with('@') {
248            // todo: create local package?
249            // if importing.starts_with("@local") { return None; }
250
251            // This is a package import, not a file import.
252            return None;
253        }
254
255        let file_id = node.span().id()?;
256        let root_path = self.ctx.path_for_id(file_id.join("/")).ok()?;
257        let path_in_workspace = file_id.vpath().join(importing.as_str());
258        let new_path = path_in_workspace.resolve(root_path.as_path())?;
259        let new_file_url = path_to_url(&new_path).ok()?;
260
261        let edit = self.create_file(new_file_url, false);
262
263        let file_to_create = unix_slash(path_in_workspace.as_rooted_path());
264        let action = CodeAction {
265            title: format!("Create missing file at `{file_to_create}`"),
266            kind: Some(CodeActionKind::QUICKFIX),
267            edit: Some(edit),
268            ..CodeAction::default()
269        };
270        self.actions.push(action);
271
272        Some(())
273    }
274
275    /// Starts to work.
276    pub fn scoped(&mut self, root: &LinkedNode, range: &Range<usize>) -> Option<()> {
277        let cursor = (range.start + 1).min(self.source.text().len());
278        let node = root.leaf_at_compat(cursor)?;
279        let mut node = &node;
280
281        let mut heading_resolved = false;
282        let mut equation_resolved = false;
283        let mut path_resolved = false;
284
285        self.wrap_actions(node, range);
286
287        loop {
288            match node.kind() {
289                // Only the deepest heading is considered
290                SyntaxKind::Heading if !heading_resolved => {
291                    heading_resolved = true;
292                    self.heading_actions(node);
293                }
294                // Only the deepest equation is considered
295                SyntaxKind::Equation if !equation_resolved => {
296                    equation_resolved = true;
297                    self.equation_actions(node);
298                }
299                SyntaxKind::Str if !path_resolved => {
300                    path_resolved = true;
301                    self.path_actions(node, cursor);
302                }
303                _ => {}
304            }
305
306            node = node.parent()?;
307        }
308    }
309
310    fn path_actions(&mut self, node: &LinkedNode, cursor: usize) -> Option<()> {
311        // We can only process the case where the import path is a string.
312        if let Some(SyntaxClass::IncludePath(path_node) | SyntaxClass::ImportPath(path_node)) =
313            classify_syntax(node.clone(), cursor)
314        {
315            let str_node = adjust_expr(path_node)?;
316            let str_ast = str_node.cast::<ast::Str>()?;
317            return self.path_rewrite(self.source.id(), &str_ast.get(), &str_node);
318        }
319
320        let link_parent = node_ancestors(node)
321            .find(|node| matches!(node.kind(), SyntaxKind::FuncCall))
322            .unwrap_or(node);
323
324        // Actually there should be only one link left
325        if let Some(link_info) = get_link_exprs_in(link_parent) {
326            let objects = link_info.objects.into_iter();
327            let object_under_node = objects.filter(|link| link.range.contains(&cursor));
328
329            let mut resolved = false;
330            for link in object_under_node {
331                if let LinkTarget::Path(id, path) = link.target {
332                    // todo: is there a link that is not a path string?
333                    resolved = self.path_rewrite(id, &path, node).is_some() || resolved;
334                }
335            }
336
337            return resolved.then_some(());
338        }
339
340        None
341    }
342
343    /// Rewrites absolute paths from/to relative paths.
344    fn path_rewrite(&mut self, id: TypstFileId, path: &str, node: &LinkedNode) -> Option<()> {
345        if !matches!(node.kind(), SyntaxKind::Str) {
346            log::warn!("bad path node kind on code action: {:?}", node.kind());
347            return None;
348        }
349
350        let path = Path::new(path);
351
352        if path.starts_with("/") {
353            // Convert absolute path to relative path
354            let cur_path = id.vpath().as_rooted_path().parent().unwrap();
355            let new_path = diff(path, cur_path)?;
356            let edit = self.edit_str(node, unix_slash(&new_path))?;
357            let action = CodeAction {
358                title: "Convert to relative path".to_string(),
359                kind: Some(CodeActionKind::REFACTOR_REWRITE),
360                edit: Some(edit),
361                ..CodeAction::default()
362            };
363            self.actions.push(action);
364        } else {
365            // Convert relative path to absolute path
366            let mut new_path = id.vpath().as_rooted_path().parent().unwrap().to_path_buf();
367            for i in path.components() {
368                match i {
369                    std::path::Component::ParentDir => {
370                        new_path.pop().then_some(())?;
371                    }
372                    std::path::Component::Normal(name) => {
373                        new_path.push(name);
374                    }
375                    _ => {}
376                }
377            }
378            let edit = self.edit_str(node, unix_slash(&new_path))?;
379            let action = CodeAction {
380                title: "Convert to absolute path".to_string(),
381                kind: Some(CodeActionKind::REFACTOR_REWRITE),
382                edit: Some(edit),
383                ..CodeAction::default()
384            };
385            self.actions.push(action);
386        }
387
388        Some(())
389    }
390
391    fn edit_str(&mut self, node: &LinkedNode, new_content: String) -> Option<EcoWorkspaceEdit> {
392        if !matches!(node.kind(), SyntaxKind::Str) {
393            log::warn!("edit_str only works on string AST nodes: {:?}", node.kind());
394            return None;
395        }
396
397        self.local_edit(EcoSnippetTextEdit::new_plain(
398            self.ctx.to_lsp_range(node.range(), &self.source),
399            // todo: this is merely occasionally correct, abusing string escape (`fmt::Debug`)
400            eco_format!("{new_content:?}"),
401        ))
402    }
403
404    fn wrap_actions(&mut self, node: &LinkedNode, range: &Range<usize>) -> Option<()> {
405        if range.is_empty() {
406            return None;
407        }
408
409        let start_mode = interpret_mode_at(Some(node));
410        if !matches!(start_mode, InterpretMode::Markup | InterpretMode::Math) {
411            return None;
412        }
413
414        let edit = self.local_edits(vec![
415            EcoSnippetTextEdit::new_plain(
416                self.ctx
417                    .to_lsp_range(range.start..range.start, &self.source),
418                EcoString::inline("#["),
419            ),
420            EcoSnippetTextEdit::new_plain(
421                self.ctx.to_lsp_range(range.end..range.end, &self.source),
422                EcoString::inline("]"),
423            ),
424        ])?;
425
426        let action = CodeAction {
427            title: "Wrap with content block".to_string(),
428            kind: Some(CodeActionKind::REFACTOR_REWRITE),
429            edit: Some(edit),
430            ..CodeAction::default()
431        };
432        self.actions.push(action);
433
434        Some(())
435    }
436
437    fn heading_actions(&mut self, node: &LinkedNode) -> Option<()> {
438        let heading = node.cast::<ast::Heading>()?;
439        let depth = heading.depth().get();
440
441        // Only the marker is replaced, for minimal text change
442        let marker = node
443            .children()
444            .find(|child| child.kind() == SyntaxKind::HeadingMarker)?;
445        let marker_range = marker.range();
446
447        if depth > 1 {
448            // Decrease depth of heading
449            let action = CodeAction {
450                title: "Decrease depth of heading".to_string(),
451                kind: Some(CodeActionKind::REFACTOR_REWRITE),
452                edit: Some(self.local_edit(EcoSnippetTextEdit::new_plain(
453                    self.ctx.to_lsp_range(marker_range.clone(), &self.source),
454                    EcoString::inline("=").repeat(depth - 1),
455                ))?),
456                ..CodeAction::default()
457            };
458            self.actions.push(action);
459        }
460
461        // Increase depth of heading
462        let action = CodeAction {
463            title: "Increase depth of heading".to_string(),
464            kind: Some(CodeActionKind::REFACTOR_REWRITE),
465            edit: Some(self.local_edit(EcoSnippetTextEdit::new_plain(
466                self.ctx.to_lsp_range(marker_range, &self.source),
467                EcoString::inline("=").repeat(depth + 1),
468            ))?),
469            ..CodeAction::default()
470        };
471        self.actions.push(action);
472
473        Some(())
474    }
475
476    fn equation_actions(&mut self, node: &LinkedNode) -> Option<()> {
477        let equation = node.cast::<ast::Equation>()?;
478        let body = equation.body();
479        let is_block = equation.block();
480
481        let body = node.find(body.span())?;
482        let body_range = body.range();
483        let node_end = node.range().end;
484
485        let mut chs = node.children();
486        let chs = chs.by_ref();
487        let is_dollar = |node: &LinkedNode| node.kind() == SyntaxKind::Dollar;
488        let first_dollar = chs.take(1).find(is_dollar)?;
489        let last_dollar = chs.rev().take(1).find(is_dollar)?;
490
491        // Erroneous equation is skipped.
492        // For example, some unclosed equation.
493        if first_dollar.offset() == last_dollar.offset() {
494            return None;
495        }
496
497        let front_range = self
498            .ctx
499            .to_lsp_range(first_dollar.range().end..body_range.start, &self.source);
500        let back_range = self
501            .ctx
502            .to_lsp_range(body_range.end..last_dollar.range().start, &self.source);
503
504        // Retrieve punctuation to move
505        let mark_after_equation = self
506            .source
507            .text()
508            .get(node_end..)
509            .and_then(|text| {
510                let mut ch = text.chars();
511                let nx = ch.next()?;
512                Some((nx, ch.next()))
513            })
514            .filter(|(ch, ch_next)| {
515                static IS_PUNCTUATION: LazyLock<Regex> =
516                    LazyLock::new(|| Regex::new(r"\p{Punctuation}").unwrap());
517                (ch.is_ascii_punctuation()
518                    && ch_next.is_none_or(|ch_next| !ch_next.is_ascii_punctuation()))
519                    || (!ch.is_ascii_punctuation() && IS_PUNCTUATION.is_match(&ch.to_string()))
520            });
521        let punc_modify = if let Some((nx, _)) = mark_after_equation {
522            let ch_range = self
523                .ctx
524                .to_lsp_range(node_end..node_end + nx.len_utf8(), &self.source);
525            let remove_edit = EcoSnippetTextEdit::new_plain(ch_range, EcoString::new());
526            Some((nx, remove_edit))
527        } else {
528            None
529        };
530
531        let rewrite_action = |title: &str, new_text: &str| {
532            let mut edits = vec![
533                EcoSnippetTextEdit::new_plain(front_range, new_text.into()),
534                EcoSnippetTextEdit::new_plain(
535                    back_range,
536                    if !new_text.is_empty() {
537                        if let Some((ch, _)) = &punc_modify {
538                            EcoString::from(*ch) + new_text
539                        } else {
540                            new_text.into()
541                        }
542                    } else {
543                        EcoString::new()
544                    },
545                ),
546            ];
547
548            if !new_text.is_empty()
549                && let Some((_, edit)) = &punc_modify
550            {
551                edits.push(edit.clone());
552            }
553
554            Some(CodeAction {
555                title: title.to_owned(),
556                kind: Some(CodeActionKind::REFACTOR_REWRITE),
557                edit: Some(self.local_edits(edits)?),
558                ..CodeAction::default()
559            })
560        };
561
562        // Prepare actions
563        let toggle_action = if is_block {
564            rewrite_action("Convert to inline equation", "")?
565        } else {
566            rewrite_action("Convert to block equation", " ")?
567        };
568        let block_action = rewrite_action("Convert to multiple-line block equation", "\n");
569
570        self.actions.push(toggle_action);
571        if let Some(a2) = block_action {
572            self.actions.push(a2);
573        }
574
575        Some(())
576    }
577
578    fn create_file(&self, uri: Url, needs_confirmation: bool) -> EcoWorkspaceEdit {
579        let change_id = "Typst Create Missing Files".to_string();
580
581        let create_op = EcoDocumentChangeOperation::Op(lsp_types::ResourceOp::Create(CreateFile {
582            uri,
583            options: Some(CreateFileOptions {
584                overwrite: Some(false),
585                ignore_if_exists: None,
586            }),
587            annotation_id: Some(change_id.clone()),
588        }));
589
590        let mut change_annotations = HashMap::new();
591        change_annotations.insert(
592            change_id.clone(),
593            ChangeAnnotation {
594                label: change_id,
595                needs_confirmation: Some(needs_confirmation),
596                description: Some("The file is missing but required by code".to_string()),
597            },
598        );
599
600        EcoWorkspaceEdit {
601            changes: None,
602            document_changes: Some(EcoDocumentChanges::Operations(vec![create_op])),
603            change_annotations: Some(change_annotations),
604        }
605    }
606}
607
608#[derive(Debug, Clone, Copy)]
609enum AutofixKind {
610    UnknownVariable,
611    FileNotFound,
612}
613
614fn match_autofix_kind(msg: &str) -> Option<AutofixKind> {
615    static PATTERNS: &[(&str, AutofixKind)] = &[
616        ("unknown variable", AutofixKind::UnknownVariable),
617        ("file not found", AutofixKind::FileNotFound),
618    ];
619
620    for (pattern, kind) in PATTERNS {
621        if msg.starts_with(pattern) {
622            return Some(*kind);
623        }
624    }
625
626    None
627}