syntax/ast/
edit.rs

1//! This module contains functions for editing syntax trees. As the trees are
2//! immutable, all function here return a fresh copy of the tree, instead of
3//! doing an in-place modification.
4use std::{fmt, iter, ops};
5
6use crate::{
7    AstToken, NodeOrToken, SyntaxElement, SyntaxNode, SyntaxToken,
8    ast::{self, AstNode, make},
9    syntax_editor::{SyntaxEditor, SyntaxMappingBuilder},
10    ted,
11};
12
13use super::syntax_factory::SyntaxFactory;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub struct IndentLevel(pub u8);
17
18impl From<u8> for IndentLevel {
19    fn from(level: u8) -> IndentLevel {
20        IndentLevel(level)
21    }
22}
23
24impl fmt::Display for IndentLevel {
25    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
26        let spaces = "                                        ";
27        let buf;
28        let len = self.0 as usize * 4;
29        let indent = if len <= spaces.len() {
30            &spaces[..len]
31        } else {
32            buf = " ".repeat(len);
33            &buf
34        };
35        fmt::Display::fmt(indent, f)
36    }
37}
38
39impl ops::Add<u8> for IndentLevel {
40    type Output = IndentLevel;
41    fn add(self, rhs: u8) -> IndentLevel {
42        IndentLevel(self.0 + rhs)
43    }
44}
45
46impl IndentLevel {
47    pub fn single() -> IndentLevel {
48        IndentLevel(0)
49    }
50    pub fn is_zero(&self) -> bool {
51        self.0 == 0
52    }
53    pub fn from_element(element: &SyntaxElement) -> IndentLevel {
54        match element {
55            rowan::NodeOrToken::Node(it) => IndentLevel::from_node(it),
56            rowan::NodeOrToken::Token(it) => IndentLevel::from_token(it),
57        }
58    }
59
60    pub fn from_node(node: &SyntaxNode) -> IndentLevel {
61        match node.first_token() {
62            Some(it) => Self::from_token(&it),
63            None => IndentLevel(0),
64        }
65    }
66
67    pub fn from_token(token: &SyntaxToken) -> IndentLevel {
68        for ws in prev_tokens(token.clone()).filter_map(ast::Whitespace::cast) {
69            let text = ws.syntax().text();
70            if let Some(pos) = text.rfind('\n') {
71                let level = text[pos + 1..].chars().count() / 4;
72                return IndentLevel(level as u8);
73            }
74        }
75        IndentLevel(0)
76    }
77
78    /// XXX: this intentionally doesn't change the indent of the very first token.
79    /// For example, in something like:
80    /// ```
81    /// fn foo() -> i32 {
82    ///    92
83    /// }
84    /// ```
85    /// if you indent the block, the `{` token would stay put.
86    pub(super) fn increase_indent(self, node: &SyntaxNode) {
87        let tokens = node.preorder_with_tokens().filter_map(|event| match event {
88            rowan::WalkEvent::Leave(NodeOrToken::Token(it)) => Some(it),
89            _ => None,
90        });
91        for token in tokens {
92            if let Some(ws) = ast::Whitespace::cast(token)
93                && ws.text().contains('\n')
94            {
95                let new_ws = make::tokens::whitespace(&format!("{}{self}", ws.syntax()));
96                ted::replace(ws.syntax(), &new_ws);
97            }
98        }
99    }
100
101    pub(super) fn clone_increase_indent(self, node: &SyntaxNode) -> SyntaxNode {
102        let node = node.clone_subtree();
103        let mut editor = SyntaxEditor::new(node.clone());
104        let tokens = node
105            .preorder_with_tokens()
106            .filter_map(|event| match event {
107                rowan::WalkEvent::Leave(NodeOrToken::Token(it)) => Some(it),
108                _ => None,
109            })
110            .filter_map(ast::Whitespace::cast)
111            .filter(|ws| ws.text().contains('\n'));
112        for ws in tokens {
113            let new_ws = make::tokens::whitespace(&format!("{}{self}", ws.syntax()));
114            editor.replace(ws.syntax(), &new_ws);
115        }
116        editor.finish().new_root().clone()
117    }
118
119    pub(super) fn decrease_indent(self, node: &SyntaxNode) {
120        let tokens = node.preorder_with_tokens().filter_map(|event| match event {
121            rowan::WalkEvent::Leave(NodeOrToken::Token(it)) => Some(it),
122            _ => None,
123        });
124        for token in tokens {
125            if let Some(ws) = ast::Whitespace::cast(token)
126                && ws.text().contains('\n')
127            {
128                let new_ws = make::tokens::whitespace(
129                    &ws.syntax().text().replace(&format!("\n{self}"), "\n"),
130                );
131                ted::replace(ws.syntax(), &new_ws);
132            }
133        }
134    }
135
136    pub(super) fn clone_decrease_indent(self, node: &SyntaxNode) -> SyntaxNode {
137        let node = node.clone_subtree();
138        let mut editor = SyntaxEditor::new(node.clone());
139        let tokens = node
140            .preorder_with_tokens()
141            .filter_map(|event| match event {
142                rowan::WalkEvent::Leave(NodeOrToken::Token(it)) => Some(it),
143                _ => None,
144            })
145            .filter_map(ast::Whitespace::cast)
146            .filter(|ws| ws.text().contains('\n'));
147        for ws in tokens {
148            let new_ws =
149                make::tokens::whitespace(&ws.syntax().text().replace(&format!("\n{self}"), "\n"));
150            editor.replace(ws.syntax(), &new_ws);
151        }
152        editor.finish().new_root().clone()
153    }
154}
155
156fn prev_tokens(token: SyntaxToken) -> impl Iterator<Item = SyntaxToken> {
157    iter::successors(Some(token), |token| token.prev_token())
158}
159
160pub trait AstNodeEdit: AstNode + Clone + Sized {
161    fn indent_level(&self) -> IndentLevel {
162        IndentLevel::from_node(self.syntax())
163    }
164    #[must_use]
165    fn indent(&self, level: IndentLevel) -> Self {
166        Self::cast(level.clone_increase_indent(self.syntax())).unwrap()
167    }
168    #[must_use]
169    fn indent_with_mapping(&self, level: IndentLevel, make: &SyntaxFactory) -> Self {
170        let new_node = self.indent(level);
171        if let Some(mut mapping) = make.mappings() {
172            let mut builder = SyntaxMappingBuilder::new(new_node.syntax().clone());
173            for (old, new) in self.syntax().children().zip(new_node.syntax().children()) {
174                builder.map_node(old, new);
175            }
176            builder.finish(&mut mapping);
177        }
178        new_node
179    }
180    #[must_use]
181    fn dedent(&self, level: IndentLevel) -> Self {
182        Self::cast(level.clone_decrease_indent(self.syntax())).unwrap()
183    }
184    #[must_use]
185    fn reset_indent(&self) -> Self {
186        let level = IndentLevel::from_node(self.syntax());
187        self.dedent(level)
188    }
189}
190
191impl<N: AstNode + Clone> AstNodeEdit for N {}
192
193#[test]
194fn test_increase_indent() {
195    let arm_list = {
196        let arm = make::match_arm(make::wildcard_pat().into(), None, make::ext::expr_unit());
197        make::match_arm_list([arm.clone(), arm])
198    };
199    assert_eq!(
200        arm_list.syntax().to_string(),
201        "{
202    _ => (),
203    _ => (),
204}"
205    );
206    let indented = arm_list.indent(IndentLevel(2));
207    assert_eq!(
208        indented.syntax().to_string(),
209        "{
210            _ => (),
211            _ => (),
212        }"
213    );
214}