1use 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
17pub struct CodeActionWorker<'a> {
19 ctx: &'a mut LocalContext,
21 source: Source,
23 pub actions: Vec<CodeAction>,
25 local_url: OnceLock<Option<Url>>,
27}
28
29impl<'a> CodeActionWorker<'a> {
30 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 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 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 fn add_spaces_to_math_unknown_variable(&mut self, node: &LinkedNode<'_>) -> Option<()> {
211 let ident = node.cast::<ast::MathIdent>()?.get();
212
213 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 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 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 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 SyntaxKind::Heading if !heading_resolved => {
291 heading_resolved = true;
292 self.heading_actions(node);
293 }
294 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 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 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 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 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 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 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 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 let marker = node
443 .children()
444 .find(|child| child.kind() == SyntaxKind::HeadingMarker)?;
445 let marker_range = marker.range();
446
447 if depth > 1 {
448 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 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 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 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 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}