syntax/ast/
edit_in_place.rs

1//! Structural editing for ast.
2
3use std::iter::{empty, once, successors};
4
5use parser::{SyntaxKind, T};
6
7use crate::{
8    AstNode, AstToken, Direction, SyntaxElement,
9    SyntaxKind::{ATTR, COMMENT, WHITESPACE},
10    SyntaxNode, SyntaxToken,
11    algo::{self, neighbor},
12    ast::{self, HasGenericParams, edit::IndentLevel, make},
13    ted::{self, Position},
14};
15
16use super::{GenericParam, HasName};
17
18pub trait GenericParamsOwnerEdit: ast::HasGenericParams {
19    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList;
20    fn get_or_create_where_clause(&self) -> ast::WhereClause;
21}
22
23impl GenericParamsOwnerEdit for ast::Fn {
24    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList {
25        match self.generic_param_list() {
26            Some(it) => it,
27            None => {
28                let position = if let Some(name) = self.name() {
29                    Position::after(name.syntax)
30                } else if let Some(fn_token) = self.fn_token() {
31                    Position::after(fn_token)
32                } else if let Some(param_list) = self.param_list() {
33                    Position::before(param_list.syntax)
34                } else {
35                    Position::last_child_of(self.syntax())
36                };
37                create_generic_param_list(position)
38            }
39        }
40    }
41
42    fn get_or_create_where_clause(&self) -> ast::WhereClause {
43        if self.where_clause().is_none() {
44            let position = if let Some(ty) = self.ret_type() {
45                Position::after(ty.syntax())
46            } else if let Some(param_list) = self.param_list() {
47                Position::after(param_list.syntax())
48            } else {
49                Position::last_child_of(self.syntax())
50            };
51            create_where_clause(position);
52        }
53        self.where_clause().unwrap()
54    }
55}
56
57impl GenericParamsOwnerEdit for ast::Impl {
58    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList {
59        match self.generic_param_list() {
60            Some(it) => it,
61            None => {
62                let position = match self.impl_token() {
63                    Some(imp_token) => Position::after(imp_token),
64                    None => Position::last_child_of(self.syntax()),
65                };
66                create_generic_param_list(position)
67            }
68        }
69    }
70
71    fn get_or_create_where_clause(&self) -> ast::WhereClause {
72        if self.where_clause().is_none() {
73            let position = match self.assoc_item_list() {
74                Some(items) => Position::before(items.syntax()),
75                None => Position::last_child_of(self.syntax()),
76            };
77            create_where_clause(position);
78        }
79        self.where_clause().unwrap()
80    }
81}
82
83impl GenericParamsOwnerEdit for ast::Trait {
84    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList {
85        match self.generic_param_list() {
86            Some(it) => it,
87            None => {
88                let position = if let Some(name) = self.name() {
89                    Position::after(name.syntax)
90                } else if let Some(trait_token) = self.trait_token() {
91                    Position::after(trait_token)
92                } else {
93                    Position::last_child_of(self.syntax())
94                };
95                create_generic_param_list(position)
96            }
97        }
98    }
99
100    fn get_or_create_where_clause(&self) -> ast::WhereClause {
101        if self.where_clause().is_none() {
102            let position = match (self.assoc_item_list(), self.semicolon_token()) {
103                (Some(items), _) => Position::before(items.syntax()),
104                (_, Some(tok)) => Position::before(tok),
105                (None, None) => Position::last_child_of(self.syntax()),
106            };
107            create_where_clause(position);
108        }
109        self.where_clause().unwrap()
110    }
111}
112
113impl GenericParamsOwnerEdit for ast::TypeAlias {
114    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList {
115        match self.generic_param_list() {
116            Some(it) => it,
117            None => {
118                let position = if let Some(name) = self.name() {
119                    Position::after(name.syntax)
120                } else if let Some(trait_token) = self.type_token() {
121                    Position::after(trait_token)
122                } else {
123                    Position::last_child_of(self.syntax())
124                };
125                create_generic_param_list(position)
126            }
127        }
128    }
129
130    fn get_or_create_where_clause(&self) -> ast::WhereClause {
131        if self.where_clause().is_none() {
132            let position = match self.eq_token() {
133                Some(tok) => Position::before(tok),
134                None => match self.semicolon_token() {
135                    Some(tok) => Position::before(tok),
136                    None => Position::last_child_of(self.syntax()),
137                },
138            };
139            create_where_clause(position);
140        }
141        self.where_clause().unwrap()
142    }
143}
144
145impl GenericParamsOwnerEdit for ast::Struct {
146    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList {
147        match self.generic_param_list() {
148            Some(it) => it,
149            None => {
150                let position = if let Some(name) = self.name() {
151                    Position::after(name.syntax)
152                } else if let Some(struct_token) = self.struct_token() {
153                    Position::after(struct_token)
154                } else {
155                    Position::last_child_of(self.syntax())
156                };
157                create_generic_param_list(position)
158            }
159        }
160    }
161
162    fn get_or_create_where_clause(&self) -> ast::WhereClause {
163        if self.where_clause().is_none() {
164            let tfl = self.field_list().and_then(|fl| match fl {
165                ast::FieldList::RecordFieldList(_) => None,
166                ast::FieldList::TupleFieldList(it) => Some(it),
167            });
168            let position = if let Some(tfl) = tfl {
169                Position::after(tfl.syntax())
170            } else if let Some(gpl) = self.generic_param_list() {
171                Position::after(gpl.syntax())
172            } else if let Some(name) = self.name() {
173                Position::after(name.syntax())
174            } else {
175                Position::last_child_of(self.syntax())
176            };
177            create_where_clause(position);
178        }
179        self.where_clause().unwrap()
180    }
181}
182
183impl GenericParamsOwnerEdit for ast::Enum {
184    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList {
185        match self.generic_param_list() {
186            Some(it) => it,
187            None => {
188                let position = if let Some(name) = self.name() {
189                    Position::after(name.syntax)
190                } else if let Some(enum_token) = self.enum_token() {
191                    Position::after(enum_token)
192                } else {
193                    Position::last_child_of(self.syntax())
194                };
195                create_generic_param_list(position)
196            }
197        }
198    }
199
200    fn get_or_create_where_clause(&self) -> ast::WhereClause {
201        if self.where_clause().is_none() {
202            let position = if let Some(gpl) = self.generic_param_list() {
203                Position::after(gpl.syntax())
204            } else if let Some(name) = self.name() {
205                Position::after(name.syntax())
206            } else {
207                Position::last_child_of(self.syntax())
208            };
209            create_where_clause(position);
210        }
211        self.where_clause().unwrap()
212    }
213}
214
215fn create_where_clause(position: Position) {
216    let where_clause = make::where_clause(empty()).clone_for_update();
217    ted::insert(position, where_clause.syntax());
218}
219
220fn create_generic_param_list(position: Position) -> ast::GenericParamList {
221    let gpl = make::generic_param_list(empty()).clone_for_update();
222    ted::insert_raw(position, gpl.syntax());
223    gpl
224}
225
226pub trait AttrsOwnerEdit: ast::HasAttrs {
227    fn remove_attrs_and_docs(&self) {
228        remove_attrs_and_docs(self.syntax());
229
230        fn remove_attrs_and_docs(node: &SyntaxNode) {
231            let mut remove_next_ws = false;
232            for child in node.children_with_tokens() {
233                match child.kind() {
234                    ATTR | COMMENT => {
235                        remove_next_ws = true;
236                        child.detach();
237                        continue;
238                    }
239                    WHITESPACE if remove_next_ws => {
240                        child.detach();
241                    }
242                    _ => (),
243                }
244                remove_next_ws = false;
245            }
246        }
247    }
248}
249
250impl<T: ast::HasAttrs> AttrsOwnerEdit for T {}
251
252impl ast::GenericParamList {
253    pub fn add_generic_param(&self, generic_param: ast::GenericParam) {
254        match self.generic_params().last() {
255            Some(last_param) => {
256                let position = Position::after(last_param.syntax());
257                let elements = vec![
258                    make::token(T![,]).into(),
259                    make::tokens::single_space().into(),
260                    generic_param.syntax().clone().into(),
261                ];
262                ted::insert_all(position, elements);
263            }
264            None => {
265                let after_l_angle = Position::after(self.l_angle_token().unwrap());
266                ted::insert(after_l_angle, generic_param.syntax());
267            }
268        }
269    }
270
271    /// Removes the existing generic param
272    pub fn remove_generic_param(&self, generic_param: ast::GenericParam) {
273        if let Some(previous) = generic_param.syntax().prev_sibling() {
274            if let Some(next_token) = previous.next_sibling_or_token() {
275                ted::remove_all(next_token..=generic_param.syntax().clone().into());
276            }
277        } else if let Some(next) = generic_param.syntax().next_sibling() {
278            if let Some(next_token) = next.prev_sibling_or_token() {
279                ted::remove_all(generic_param.syntax().clone().into()..=next_token);
280            }
281        } else {
282            ted::remove(generic_param.syntax());
283        }
284    }
285
286    /// Find the params corresponded to generic arg
287    pub fn find_generic_arg(&self, generic_arg: &ast::GenericArg) -> Option<GenericParam> {
288        self.generic_params().find_map(move |param| match (&param, &generic_arg) {
289            (ast::GenericParam::LifetimeParam(a), ast::GenericArg::LifetimeArg(b)) => {
290                (a.lifetime()?.lifetime_ident_token()?.text()
291                    == b.lifetime()?.lifetime_ident_token()?.text())
292                .then_some(param)
293            }
294            (ast::GenericParam::TypeParam(a), ast::GenericArg::TypeArg(b)) => {
295                debug_assert_eq!(b.syntax().first_token(), b.syntax().last_token());
296                (a.name()?.text() == b.syntax().first_token()?.text()).then_some(param)
297            }
298            (ast::GenericParam::ConstParam(a), ast::GenericArg::TypeArg(b)) => {
299                debug_assert_eq!(b.syntax().first_token(), b.syntax().last_token());
300                (a.name()?.text() == b.syntax().first_token()?.text()).then_some(param)
301            }
302            _ => None,
303        })
304    }
305
306    /// Removes the corresponding generic arg
307    pub fn remove_generic_arg(&self, generic_arg: &ast::GenericArg) {
308        let param_to_remove = self.find_generic_arg(generic_arg);
309
310        if let Some(param) = &param_to_remove {
311            self.remove_generic_param(param.clone());
312        }
313    }
314
315    /// Constructs a matching [`ast::GenericArgList`]
316    pub fn to_generic_args(&self) -> ast::GenericArgList {
317        let args = self.generic_params().filter_map(|param| match param {
318            ast::GenericParam::LifetimeParam(it) => {
319                Some(ast::GenericArg::LifetimeArg(make::lifetime_arg(it.lifetime()?)))
320            }
321            ast::GenericParam::TypeParam(it) => {
322                Some(ast::GenericArg::TypeArg(make::type_arg(make::ext::ty_name(it.name()?))))
323            }
324            ast::GenericParam::ConstParam(it) => {
325                // Name-only const params get parsed as `TypeArg`s
326                Some(ast::GenericArg::TypeArg(make::type_arg(make::ext::ty_name(it.name()?))))
327            }
328        });
329
330        make::generic_arg_list(args)
331    }
332}
333
334impl ast::WhereClause {
335    pub fn add_predicate(&self, predicate: ast::WherePred) {
336        if let Some(pred) = self.predicates().last()
337            && !pred.syntax().siblings_with_tokens(Direction::Next).any(|it| it.kind() == T![,])
338        {
339            ted::append_child_raw(self.syntax(), make::token(T![,]));
340        }
341        ted::append_child(self.syntax(), predicate.syntax());
342    }
343
344    pub fn remove_predicate(&self, predicate: ast::WherePred) {
345        if let Some(previous) = predicate.syntax().prev_sibling() {
346            if let Some(next_token) = previous.next_sibling_or_token() {
347                ted::remove_all(next_token..=predicate.syntax().clone().into());
348            }
349        } else if let Some(next) = predicate.syntax().next_sibling() {
350            if let Some(next_token) = next.prev_sibling_or_token() {
351                ted::remove_all(predicate.syntax().clone().into()..=next_token);
352            }
353        } else {
354            ted::remove(predicate.syntax());
355        }
356    }
357}
358
359pub trait Removable: AstNode {
360    fn remove(&self);
361}
362
363impl Removable for ast::TypeBoundList {
364    fn remove(&self) {
365        match self.syntax().siblings_with_tokens(Direction::Prev).find(|it| it.kind() == T![:]) {
366            Some(colon) => ted::remove_all(colon..=self.syntax().clone().into()),
367            None => ted::remove(self.syntax()),
368        }
369    }
370}
371
372impl Removable for ast::UseTree {
373    fn remove(&self) {
374        for dir in [Direction::Next, Direction::Prev] {
375            if let Some(next_use_tree) = neighbor(self, dir) {
376                let separators = self
377                    .syntax()
378                    .siblings_with_tokens(dir)
379                    .skip(1)
380                    .take_while(|it| it.as_node() != Some(next_use_tree.syntax()));
381                ted::remove_all_iter(separators);
382                break;
383            }
384        }
385        ted::remove(self.syntax());
386    }
387}
388
389impl ast::UseTree {
390    /// Deletes the usetree node represented by the input. Recursively removes parents, including use nodes that become empty.
391    pub fn remove_recursive(self) {
392        let parent = self.syntax().parent();
393
394        self.remove();
395
396        if let Some(u) = parent.clone().and_then(ast::Use::cast) {
397            if u.use_tree().is_none() {
398                u.remove();
399            }
400        } else if let Some(u) = parent.and_then(ast::UseTreeList::cast) {
401            if u.use_trees().next().is_none() {
402                let parent = u.syntax().parent().and_then(ast::UseTree::cast);
403                if let Some(u) = parent {
404                    u.remove_recursive();
405                }
406            }
407            u.remove_unnecessary_braces();
408        }
409    }
410
411    pub fn get_or_create_use_tree_list(&self) -> ast::UseTreeList {
412        match self.use_tree_list() {
413            Some(it) => it,
414            None => {
415                let position = Position::last_child_of(self.syntax());
416                let use_tree_list = make::use_tree_list(empty()).clone_for_update();
417                let mut elements = Vec::with_capacity(2);
418                if self.coloncolon_token().is_none() {
419                    elements.push(make::token(T![::]).into());
420                }
421                elements.push(use_tree_list.syntax().clone().into());
422                ted::insert_all_raw(position, elements);
423                use_tree_list
424            }
425        }
426    }
427
428    /// Splits off the given prefix, making it the path component of the use tree,
429    /// appending the rest of the path to all UseTreeList items.
430    ///
431    /// # Examples
432    ///
433    /// `prefix$0::suffix` -> `prefix::{suffix}`
434    ///
435    /// `prefix$0` -> `prefix::{self}`
436    ///
437    /// `prefix$0::*` -> `prefix::{*}`
438    pub fn split_prefix(&self, prefix: &ast::Path) {
439        debug_assert_eq!(self.path(), Some(prefix.top_path()));
440        let path = self.path().unwrap();
441        if &path == prefix && self.use_tree_list().is_none() {
442            if self.star_token().is_some() {
443                // path$0::* -> *
444                if let Some(a) = self.coloncolon_token() {
445                    ted::remove(a)
446                }
447                ted::remove(prefix.syntax());
448            } else {
449                // path$0 -> self
450                let self_suffix =
451                    make::path_unqualified(make::path_segment_self()).clone_for_update();
452                ted::replace(path.syntax(), self_suffix.syntax());
453            }
454        } else if split_path_prefix(prefix).is_none() {
455            return;
456        }
457        // At this point, prefix path is detached; _self_ use tree has suffix path.
458        // Next, transform 'suffix' use tree into 'prefix::{suffix}'
459        let subtree = self.clone_subtree().clone_for_update();
460        ted::remove_all_iter(self.syntax().children_with_tokens());
461        ted::insert(Position::first_child_of(self.syntax()), prefix.syntax());
462        self.get_or_create_use_tree_list().add_use_tree(subtree);
463
464        fn split_path_prefix(prefix: &ast::Path) -> Option<()> {
465            let parent = prefix.parent_path()?;
466            let segment = parent.segment()?;
467            if algo::has_errors(segment.syntax()) {
468                return None;
469            }
470            for p in successors(parent.parent_path(), |it| it.parent_path()) {
471                p.segment()?;
472            }
473            if let Some(a) = prefix.parent_path().and_then(|p| p.coloncolon_token()) {
474                ted::remove(a)
475            }
476            ted::remove(prefix.syntax());
477            Some(())
478        }
479    }
480
481    /// Wraps the use tree in use tree list with no top level path (if it isn't already).
482    ///
483    /// # Examples
484    ///
485    /// `foo::bar` -> `{foo::bar}`
486    ///
487    /// `{foo::bar}` -> `{foo::bar}`
488    pub fn wrap_in_tree_list(&self) -> Option<()> {
489        if self.use_tree_list().is_some()
490            && self.path().is_none()
491            && self.star_token().is_none()
492            && self.rename().is_none()
493        {
494            return None;
495        }
496        let subtree = self.clone_subtree().clone_for_update();
497        ted::remove_all_iter(self.syntax().children_with_tokens());
498        ted::append_child(
499            self.syntax(),
500            make::use_tree_list(once(subtree)).clone_for_update().syntax(),
501        );
502        Some(())
503    }
504}
505
506impl ast::UseTreeList {
507    pub fn add_use_tree(&self, use_tree: ast::UseTree) {
508        let (position, elements) = match self.use_trees().last() {
509            Some(last_tree) => (
510                Position::after(last_tree.syntax()),
511                vec![
512                    make::token(T![,]).into(),
513                    make::tokens::single_space().into(),
514                    use_tree.syntax.into(),
515                ],
516            ),
517            None => {
518                let position = match self.l_curly_token() {
519                    Some(l_curly) => Position::after(l_curly),
520                    None => Position::last_child_of(self.syntax()),
521                };
522                (position, vec![use_tree.syntax.into()])
523            }
524        };
525        ted::insert_all_raw(position, elements);
526    }
527}
528
529impl Removable for ast::Use {
530    fn remove(&self) {
531        let next_ws = self
532            .syntax()
533            .next_sibling_or_token()
534            .and_then(|it| it.into_token())
535            .and_then(ast::Whitespace::cast);
536        if let Some(next_ws) = next_ws {
537            let ws_text = next_ws.syntax().text();
538            if let Some(rest) = ws_text.strip_prefix('\n') {
539                if rest.is_empty() {
540                    ted::remove(next_ws.syntax());
541                } else {
542                    ted::replace(next_ws.syntax(), make::tokens::whitespace(rest));
543                }
544            }
545        }
546        let prev_ws = self
547            .syntax()
548            .prev_sibling_or_token()
549            .and_then(|it| it.into_token())
550            .and_then(ast::Whitespace::cast);
551        if let Some(prev_ws) = prev_ws {
552            let ws_text = prev_ws.syntax().text();
553            let prev_newline = ws_text.rfind('\n').map(|x| x + 1).unwrap_or(0);
554            let rest = &ws_text[0..prev_newline];
555            if rest.is_empty() {
556                ted::remove(prev_ws.syntax());
557            } else {
558                ted::replace(prev_ws.syntax(), make::tokens::whitespace(rest));
559            }
560        }
561
562        ted::remove(self.syntax());
563    }
564}
565
566impl ast::Impl {
567    pub fn get_or_create_assoc_item_list(&self) -> ast::AssocItemList {
568        if self.assoc_item_list().is_none() {
569            let assoc_item_list = make::assoc_item_list(None).clone_for_update();
570            ted::append_child(self.syntax(), assoc_item_list.syntax());
571        }
572        self.assoc_item_list().unwrap()
573    }
574}
575
576impl ast::AssocItemList {
577    /// Adds a new associated item after all of the existing associated items.
578    ///
579    /// Attention! This function does align the first line of `item` with respect to `self`,
580    /// but it does _not_ change indentation of other lines (if any).
581    pub fn add_item(&self, item: ast::AssocItem) {
582        let (indent, position, whitespace) = match self.assoc_items().last() {
583            Some(last_item) => (
584                IndentLevel::from_node(last_item.syntax()),
585                Position::after(last_item.syntax()),
586                "\n\n",
587            ),
588            None => match self.l_curly_token() {
589                Some(l_curly) => {
590                    normalize_ws_between_braces(self.syntax());
591                    (IndentLevel::from_token(&l_curly) + 1, Position::after(&l_curly), "\n")
592                }
593                None => (IndentLevel::single(), Position::last_child_of(self.syntax()), "\n"),
594            },
595        };
596        let elements: Vec<SyntaxElement> = vec![
597            make::tokens::whitespace(&format!("{whitespace}{indent}")).into(),
598            item.syntax().clone().into(),
599        ];
600        ted::insert_all(position, elements);
601    }
602}
603
604impl ast::RecordExprFieldList {
605    pub fn add_field(&self, field: ast::RecordExprField) {
606        let is_multiline = self.syntax().text().contains_char('\n');
607        let whitespace = if is_multiline {
608            let indent = IndentLevel::from_node(self.syntax()) + 1;
609            make::tokens::whitespace(&format!("\n{indent}"))
610        } else {
611            make::tokens::single_space()
612        };
613
614        if is_multiline {
615            normalize_ws_between_braces(self.syntax());
616        }
617
618        let position = match self.fields().last() {
619            Some(last_field) => {
620                let comma = get_or_insert_comma_after(last_field.syntax());
621                Position::after(comma)
622            }
623            None => match self.l_curly_token() {
624                Some(it) => Position::after(it),
625                None => Position::last_child_of(self.syntax()),
626            },
627        };
628
629        ted::insert_all(position, vec![whitespace.into(), field.syntax().clone().into()]);
630        if is_multiline {
631            ted::insert(Position::after(field.syntax()), ast::make::token(T![,]));
632        }
633    }
634}
635
636impl ast::RecordExprField {
637    /// This will either replace the initializer, or in the case that this is a shorthand convert
638    /// the initializer into the name ref and insert the expr as the new initializer.
639    pub fn replace_expr(&self, expr: ast::Expr) {
640        if self.name_ref().is_some() {
641            match self.expr() {
642                Some(prev) => ted::replace(prev.syntax(), expr.syntax()),
643                None => ted::append_child(self.syntax(), expr.syntax()),
644            }
645            return;
646        }
647        // this is a shorthand
648        if let Some(ast::Expr::PathExpr(path_expr)) = self.expr()
649            && let Some(path) = path_expr.path()
650            && let Some(name_ref) = path.as_single_name_ref()
651        {
652            path_expr.syntax().detach();
653            let children = vec![
654                name_ref.syntax().clone().into(),
655                ast::make::token(T![:]).into(),
656                ast::make::tokens::single_space().into(),
657                expr.syntax().clone().into(),
658            ];
659            ted::insert_all_raw(Position::last_child_of(self.syntax()), children);
660        }
661    }
662}
663
664impl ast::RecordPatFieldList {
665    pub fn add_field(&self, field: ast::RecordPatField) {
666        let is_multiline = self.syntax().text().contains_char('\n');
667        let whitespace = if is_multiline {
668            let indent = IndentLevel::from_node(self.syntax()) + 1;
669            make::tokens::whitespace(&format!("\n{indent}"))
670        } else {
671            make::tokens::single_space()
672        };
673
674        if is_multiline {
675            normalize_ws_between_braces(self.syntax());
676        }
677
678        let position = match self.fields().last() {
679            Some(last_field) => {
680                let syntax = last_field.syntax();
681                let comma = get_or_insert_comma_after(syntax);
682                Position::after(comma)
683            }
684            None => match self.l_curly_token() {
685                Some(it) => Position::after(it),
686                None => Position::last_child_of(self.syntax()),
687            },
688        };
689
690        ted::insert_all(position, vec![whitespace.into(), field.syntax().clone().into()]);
691        if is_multiline {
692            ted::insert(Position::after(field.syntax()), ast::make::token(T![,]));
693        }
694    }
695}
696
697fn get_or_insert_comma_after(syntax: &SyntaxNode) -> SyntaxToken {
698    match syntax
699        .siblings_with_tokens(Direction::Next)
700        .filter_map(|it| it.into_token())
701        .find(|it| it.kind() == T![,])
702    {
703        Some(it) => it,
704        None => {
705            let comma = ast::make::token(T![,]);
706            ted::insert(Position::after(syntax), &comma);
707            comma
708        }
709    }
710}
711
712fn normalize_ws_between_braces(node: &SyntaxNode) -> Option<()> {
713    let l = node
714        .children_with_tokens()
715        .filter_map(|it| it.into_token())
716        .find(|it| it.kind() == T!['{'])?;
717    let r = node
718        .children_with_tokens()
719        .filter_map(|it| it.into_token())
720        .find(|it| it.kind() == T!['}'])?;
721
722    let indent = IndentLevel::from_node(node);
723
724    match l.next_sibling_or_token() {
725        Some(ws) if ws.kind() == SyntaxKind::WHITESPACE => {
726            if ws.next_sibling_or_token()?.into_token()? == r {
727                ted::replace(ws, make::tokens::whitespace(&format!("\n{indent}")));
728            }
729        }
730        Some(ws) if ws.kind() == T!['}'] => {
731            ted::insert(Position::after(l), make::tokens::whitespace(&format!("\n{indent}")));
732        }
733        _ => (),
734    }
735    Some(())
736}
737
738impl ast::IdentPat {
739    pub fn set_pat(&self, pat: Option<ast::Pat>) {
740        match pat {
741            None => {
742                if let Some(at_token) = self.at_token() {
743                    // Remove `@ Pat`
744                    let start = at_token.clone().into();
745                    let end = self
746                        .pat()
747                        .map(|it| it.syntax().clone().into())
748                        .unwrap_or_else(|| at_token.into());
749
750                    ted::remove_all(start..=end);
751
752                    // Remove any trailing ws
753                    if let Some(last) =
754                        self.syntax().last_token().filter(|it| it.kind() == WHITESPACE)
755                    {
756                        last.detach();
757                    }
758                }
759            }
760            Some(pat) => {
761                if let Some(old_pat) = self.pat() {
762                    // Replace existing pattern
763                    ted::replace(old_pat.syntax(), pat.syntax())
764                } else if let Some(at_token) = self.at_token() {
765                    // Have an `@` token but not a pattern yet
766                    ted::insert(ted::Position::after(at_token), pat.syntax());
767                } else {
768                    // Don't have an `@`, should have a name
769                    let name = self.name().unwrap();
770
771                    ted::insert_all(
772                        ted::Position::after(name.syntax()),
773                        vec![
774                            make::token(T![@]).into(),
775                            make::tokens::single_space().into(),
776                            pat.syntax().clone().into(),
777                        ],
778                    )
779                }
780            }
781        }
782    }
783}
784
785pub trait HasVisibilityEdit: ast::HasVisibility {
786    fn set_visibility(&self, visibility: Option<ast::Visibility>) {
787        if let Some(visibility) = visibility {
788            match self.visibility() {
789                Some(current_visibility) => {
790                    ted::replace(current_visibility.syntax(), visibility.syntax())
791                }
792                None => {
793                    let vis_before = self
794                        .syntax()
795                        .children_with_tokens()
796                        .find(|it| !matches!(it.kind(), WHITESPACE | COMMENT | ATTR))
797                        .unwrap_or_else(|| self.syntax().first_child_or_token().unwrap());
798
799                    ted::insert(ted::Position::before(vis_before), visibility.syntax());
800                }
801            }
802        } else if let Some(visibility) = self.visibility() {
803            ted::remove(visibility.syntax());
804        }
805    }
806}
807
808impl<T: ast::HasVisibility> HasVisibilityEdit for T {}
809
810pub trait Indent: AstNode + Clone + Sized {
811    fn indent_level(&self) -> IndentLevel {
812        IndentLevel::from_node(self.syntax())
813    }
814    fn indent(&self, by: IndentLevel) {
815        by.increase_indent(self.syntax());
816    }
817    fn dedent(&self, by: IndentLevel) {
818        by.decrease_indent(self.syntax());
819    }
820    fn reindent_to(&self, target_level: IndentLevel) {
821        let current_level = IndentLevel::from_node(self.syntax());
822        self.dedent(current_level);
823        self.indent(target_level);
824    }
825}
826
827impl<N: AstNode + Clone> Indent for N {}
828
829#[cfg(test)]
830mod tests {
831    use std::fmt;
832
833    use parser::Edition;
834
835    use crate::SourceFile;
836
837    use super::*;
838
839    fn ast_mut_from_text<N: AstNode>(text: &str) -> N {
840        let parse = SourceFile::parse(text, Edition::CURRENT);
841        parse.tree().syntax().descendants().find_map(N::cast).unwrap().clone_for_update()
842    }
843
844    #[test]
845    fn test_create_generic_param_list() {
846        fn check_create_gpl<N: GenericParamsOwnerEdit + fmt::Display>(before: &str, after: &str) {
847            let gpl_owner = ast_mut_from_text::<N>(before);
848            gpl_owner.get_or_create_generic_param_list();
849            assert_eq!(gpl_owner.to_string(), after);
850        }
851
852        check_create_gpl::<ast::Fn>("fn foo", "fn foo<>");
853        check_create_gpl::<ast::Fn>("fn foo() {}", "fn foo<>() {}");
854
855        check_create_gpl::<ast::Impl>("impl", "impl<>");
856        check_create_gpl::<ast::Impl>("impl Struct {}", "impl<> Struct {}");
857        check_create_gpl::<ast::Impl>("impl Trait for Struct {}", "impl<> Trait for Struct {}");
858
859        check_create_gpl::<ast::Trait>("trait Trait<>", "trait Trait<>");
860        check_create_gpl::<ast::Trait>("trait Trait<> {}", "trait Trait<> {}");
861
862        check_create_gpl::<ast::Struct>("struct A", "struct A<>");
863        check_create_gpl::<ast::Struct>("struct A;", "struct A<>;");
864        check_create_gpl::<ast::Struct>("struct A();", "struct A<>();");
865        check_create_gpl::<ast::Struct>("struct A {}", "struct A<> {}");
866
867        check_create_gpl::<ast::Enum>("enum E", "enum E<>");
868        check_create_gpl::<ast::Enum>("enum E {", "enum E<> {");
869    }
870
871    #[test]
872    fn test_increase_indent() {
873        let arm_list = ast_mut_from_text::<ast::Fn>(
874            "fn foo() {
875    ;
876    ;
877}",
878        );
879        arm_list.indent(IndentLevel(2));
880        assert_eq!(
881            arm_list.to_string(),
882            "fn foo() {
883            ;
884            ;
885        }",
886        );
887    }
888
889    #[test]
890    fn test_ident_pat_set_pat() {
891        #[track_caller]
892        fn check(before: &str, expected: &str, pat: Option<ast::Pat>) {
893            let pat = pat.map(|it| it.clone_for_update());
894
895            let ident_pat = ast_mut_from_text::<ast::IdentPat>(&format!("fn f() {{ {before} }}"));
896            ident_pat.set_pat(pat);
897
898            let after = ast_mut_from_text::<ast::IdentPat>(&format!("fn f() {{ {expected} }}"));
899            assert_eq!(ident_pat.to_string(), after.to_string());
900        }
901
902        // replacing
903        check("let a @ _;", "let a @ ();", Some(make::tuple_pat([]).into()));
904
905        // note: no trailing semicolon is added for the below tests since it
906        // seems to be picked up by the ident pat during error recovery?
907
908        // adding
909        check("let a ", "let a @ ()", Some(make::tuple_pat([]).into()));
910        check("let a @ ", "let a @ ()", Some(make::tuple_pat([]).into()));
911
912        // removing
913        check("let a @ ()", "let a", None);
914        check("let a @ ", "let a", None);
915    }
916}
syntax/ast/edit_in_place.rs

syntax/ast/
edit_in_place.rs
