1 //! Utilities for formatting macro expanded nodes until we get a proper formatter.
2 use syntax::{
3 ast::make,
4 ted::{self, Position},
5 NodeOrToken,
6 SyntaxKind::{self, *},
7 SyntaxNode, SyntaxToken, WalkEvent, T,
8 };
9
10 // FIXME: It would also be cool to share logic here and in the mbe tests,
11 // which are pretty unreadable at the moment.
12 /// Renders a [`SyntaxNode`] with whitespace inserted between tokens that require them.
insert_ws_into(syn: SyntaxNode) -> SyntaxNode13 pub fn insert_ws_into(syn: SyntaxNode) -> SyntaxNode {
14 let mut indent = 0;
15 let mut last: Option<SyntaxKind> = None;
16 let mut mods = Vec::new();
17 let syn = syn.clone_subtree().clone_for_update();
18
19 let before = Position::before;
20 let after = Position::after;
21
22 let do_indent = |pos: fn(_) -> Position, token: &SyntaxToken, indent| {
23 (pos(token.clone()), make::tokens::whitespace(&" ".repeat(2 * indent)))
24 };
25 let do_ws = |pos: fn(_) -> Position, token: &SyntaxToken| {
26 (pos(token.clone()), make::tokens::single_space())
27 };
28 let do_nl = |pos: fn(_) -> Position, token: &SyntaxToken| {
29 (pos(token.clone()), make::tokens::single_newline())
30 };
31
32 for event in syn.preorder_with_tokens() {
33 let token = match event {
34 WalkEvent::Enter(NodeOrToken::Token(token)) => token,
35 WalkEvent::Leave(NodeOrToken::Node(node))
36 if matches!(
37 node.kind(),
38 ATTR | MATCH_ARM | STRUCT | ENUM | UNION | FN | IMPL | MACRO_RULES
39 ) =>
40 {
41 if indent > 0 {
42 mods.push((
43 Position::after(node.clone()),
44 make::tokens::whitespace(&" ".repeat(2 * indent)),
45 ));
46 }
47 if node.parent().is_some() {
48 mods.push((Position::after(node), make::tokens::single_newline()));
49 }
50 continue;
51 }
52 _ => continue,
53 };
54 let tok = &token;
55
56 let is_next = |f: fn(SyntaxKind) -> bool, default| -> bool {
57 tok.next_token().map(|it| f(it.kind())).unwrap_or(default)
58 };
59 let is_last =
60 |f: fn(SyntaxKind) -> bool, default| -> bool { last.map(f).unwrap_or(default) };
61
62 match tok.kind() {
63 k if is_text(k)
64 && is_next(|it| !it.is_punct() || matches!(it, T![_] | T![#]), false) =>
65 {
66 mods.push(do_ws(after, tok));
67 }
68 L_CURLY if is_next(|it| it != R_CURLY, true) => {
69 indent += 1;
70 if is_last(is_text, false) {
71 mods.push(do_ws(before, tok));
72 }
73
74 mods.push(do_indent(after, tok, indent));
75 mods.push(do_nl(after, tok));
76 }
77 R_CURLY if is_last(|it| it != L_CURLY, true) => {
78 indent = indent.saturating_sub(1);
79
80 if indent > 0 {
81 mods.push(do_indent(before, tok, indent));
82 }
83 mods.push(do_nl(before, tok));
84 }
85 R_CURLY => {
86 if indent > 0 {
87 mods.push(do_indent(after, tok, indent));
88 }
89 mods.push(do_nl(after, tok));
90 }
91 LIFETIME_IDENT if is_next(is_text, true) => {
92 mods.push(do_ws(after, tok));
93 }
94 MUT_KW if is_next(|it| it == SELF_KW, false) => {
95 mods.push(do_ws(after, tok));
96 }
97 AS_KW | DYN_KW | IMPL_KW | CONST_KW => {
98 mods.push(do_ws(after, tok));
99 }
100 T![;] if is_next(|it| it != R_CURLY, true) => {
101 if indent > 0 {
102 mods.push(do_indent(after, tok, indent));
103 }
104 mods.push(do_nl(after, tok));
105 }
106 T![=] if is_next(|it| it == T![>], false) => {
107 // FIXME: this branch is for `=>` in macro_rules!, which is currently parsed as
108 // two separate symbols.
109 mods.push(do_ws(before, tok));
110 mods.push(do_ws(after, &tok.next_token().unwrap()));
111 }
112 T![->] | T![=] | T![=>] => {
113 mods.push(do_ws(before, tok));
114 mods.push(do_ws(after, tok));
115 }
116 T![!] if is_last(|it| it == MACRO_RULES_KW, false) && is_next(is_text, false) => {
117 mods.push(do_ws(after, tok));
118 }
119 _ => (),
120 }
121
122 last = Some(tok.kind());
123 }
124
125 for (pos, insert) in mods {
126 ted::insert(pos, insert);
127 }
128
129 if let Some(it) = syn.last_token().filter(|it| it.kind() == SyntaxKind::WHITESPACE) {
130 ted::remove(it);
131 }
132
133 syn
134 }
135
is_text(k: SyntaxKind) -> bool136 fn is_text(k: SyntaxKind) -> bool {
137 k.is_keyword() || k.is_literal() || k == IDENT || k == UNDERSCORE
138 }
139