1pub mod nodes;
8pub mod tokens;
9
10pub use nodes::{
11 Begin, Command, Conditional, ConditionalBranch, End, Environment, Group, NameGroup, Optional,
12};
13pub use tokens::ControlWord;
14
15use rowan::{NodeOrToken, TextRange};
16
17use crate::syntax::{SyntaxKind, SyntaxNode, SyntaxToken};
18
19pub trait AstNode {
21 fn can_cast(kind: SyntaxKind) -> bool
22 where
23 Self: Sized;
24 fn cast(syntax: SyntaxNode) -> Option<Self>
25 where
26 Self: Sized;
27 fn syntax(&self) -> &SyntaxNode;
28}
29
30pub trait AstToken {
32 fn can_cast(kind: SyntaxKind) -> bool
33 where
34 Self: Sized;
35 fn cast(syntax: SyntaxToken) -> Option<Self>
36 where
37 Self: Sized;
38 fn syntax(&self) -> &SyntaxToken;
39 fn text(&self) -> &str {
40 self.syntax().text()
41 }
42}
43
44pub fn child<N: AstNode>(parent: &SyntaxNode) -> Option<N> {
46 parent.children().find_map(N::cast)
47}
48
49pub fn children<N: AstNode>(parent: &SyntaxNode) -> impl Iterator<Item = N> {
51 parent.children().filter_map(N::cast)
52}
53
54pub fn child_token<T: AstToken>(parent: &SyntaxNode) -> Option<T> {
56 parent
57 .children_with_tokens()
58 .filter_map(NodeOrToken::into_token)
59 .find_map(T::cast)
60}
61
62pub fn command_name(command: &SyntaxNode) -> Option<String> {
74 child_token::<ControlWord>(command).map(|cw| cw.name())
75}
76
77pub fn control_word_range(command: &SyntaxNode) -> Option<TextRange> {
80 child_token::<ControlWord>(command).map(|cw| cw.range())
81}
82
83pub fn nth_group_text(command: &SyntaxNode, n: usize) -> Option<String> {
85 children::<Group>(command)
86 .nth(n)
87 .and_then(|g| g.inner_text())
88}
89
90pub fn nth_group_inner(command: &SyntaxNode, n: usize) -> Option<(TextRange, String)> {
93 children::<Group>(command).nth(n).and_then(|g| g.inner())
94}
95
96pub fn nth_group(command: &SyntaxNode, n: usize) -> Option<SyntaxNode> {
98 children::<Group>(command)
99 .nth(n)
100 .map(|g| g.syntax().clone())
101}
102
103pub fn first_group_range(command: &SyntaxNode) -> TextRange {
106 match children::<Group>(command).next() {
107 Some(group) => TextRange::new(
108 command.text_range().start(),
109 group.syntax().text_range().end(),
110 ),
111 None => command.text_range(),
112 }
113}
114
115pub fn group_command_name(group: &SyntaxNode) -> Option<String> {
120 if let Some(name) = child::<Command>(group).and_then(|c| c.name()) {
121 return Some(name);
122 }
123 let verb = group
124 .children_with_tokens()
125 .filter_map(|e| e.into_token())
126 .find(|t| t.kind() == SyntaxKind::VERB)?;
127 let name = verb.text().trim().strip_prefix('\\')?;
128 (!name.is_empty()
129 && name
130 .chars()
131 .all(|c| c.is_alphanumeric() || c == '@' || c == '_' || c == ':'))
132 .then(|| name.to_owned())
133}
134
135pub fn group_inner_source(group: &SyntaxNode) -> String {
137 nodes::inner_source_of(group)
138}
139
140pub fn environment_name(begin_or_end: &SyntaxNode) -> Option<String> {
147 if begin_or_end.kind() == SyntaxKind::BEGIN {
148 return Begin::cast(begin_or_end.clone())?.name();
149 }
150 child::<NameGroup>(begin_or_end).and_then(|g| g.text())
151}
152
153pub fn environment_name_range(begin_or_end: &SyntaxNode) -> Option<TextRange> {
156 child::<NameGroup>(begin_or_end).and_then(|g| g.range())
157}
158
159#[cfg(test)]
160mod tests {
161 use super::*;
162 use crate::parser::parse;
163
164 fn command(src: &str) -> SyntaxNode {
165 SyntaxNode::new_root(parse(src).green)
166 .descendants()
167 .find(|node| node.kind() == SyntaxKind::COMMAND)
168 .expect("a COMMAND node")
169 }
170
171 fn node(src: &str, kind: SyntaxKind) -> SyntaxNode {
172 SyntaxNode::new_root(parse(src).green)
173 .descendants()
174 .find(|n| n.kind() == kind)
175 .expect("a matching node")
176 }
177
178 #[test]
179 fn command_name_strips_backslash() {
180 assert_eq!(
181 command_name(&command("\\section{Hi}\n")).as_deref(),
182 Some("section")
183 );
184 }
185
186 #[test]
187 fn nth_group_text_reassembles_inner_tokens() {
188 assert_eq!(
189 nth_group_text(&command("\\label{sec:intro}\n"), 0).as_deref(),
190 Some("sec:intro")
191 );
192 }
193
194 #[test]
195 fn nth_group_inner_spans_only_the_key() {
196 let src = "\\label{sec:intro}\n";
198 let cmd = command(src);
199 let (range, text) = nth_group_inner(&cmd, 0).expect("an inner span");
200 assert_eq!(text, "sec:intro");
201 assert_eq!(&src[range], "sec:intro");
202 }
203
204 #[test]
205 fn nth_group_inner_empty_group_is_zero_width_after_brace() {
206 let cmd = command("\\label{}\n");
207 let (range, text) = nth_group_inner(&cmd, 0).expect("an inner span");
208 assert!(text.is_empty());
209 assert!(range.is_empty());
210 }
211
212 #[test]
213 fn nth_group_inner_none_for_nested_command() {
214 assert_eq!(nth_group_inner(&command("\\input{\\jobname}\n"), 0), None);
215 }
216
217 #[test]
218 fn nth_group_inner_none_for_parameter_token() {
219 assert_eq!(nth_group_inner(&command("\\ref{#1}\n"), 0), None);
222 assert_eq!(nth_group_inner(&command("\\eqref{##1}\n"), 0), None);
223 assert_eq!(nth_group_text(&command("\\input{#1}\n"), 0), None);
224 }
225
226 #[test]
227 fn nth_group_text_none_for_nested_command() {
228 assert_eq!(nth_group_text(&command("\\input{\\jobname}\n"), 0), None);
229 }
230
231 #[test]
232 fn nth_group_text_none_when_group_absent() {
233 assert_eq!(nth_group_text(&command("\\input\n"), 0), None);
234 }
235
236 #[test]
237 fn group_command_name_reads_braced_control_word() {
238 let cmd = command("\\newcommand{\\foo}{x}\n");
239 let name = nth_group(&cmd, 0).and_then(|g| group_command_name(&g));
240 assert_eq!(name.as_deref(), Some("foo"));
241 }
242
243 #[test]
244 fn group_command_name_none_for_plain_text() {
245 let cmd = command("\\newenvironment{thm}{a}{b}\n");
246 let name = nth_group(&cmd, 0).and_then(|g| group_command_name(&g));
247 assert_eq!(name, None);
248 }
249
250 #[test]
251 fn group_inner_source_keeps_nested_braces() {
252 let cmd = command("\\NewDocumentCommand{\\foo}{m O{d} m}{x}\n");
255 let spec = nth_group(&cmd, 1).map(|g| group_inner_source(&g));
256 assert_eq!(spec.as_deref(), Some("m O{d} m"));
257 assert_eq!(nth_group_text(&cmd, 1), None);
258 }
259
260 #[test]
261 fn environment_name_range_spans_only_the_name() {
262 let src = "\\begin{equation}\nx\n\\end{equation}\n";
263 let begin = node(src, SyntaxKind::BEGIN);
264 let range = environment_name_range(&begin).expect("a name span");
265 assert_eq!(&src[range], "equation");
266
267 let end = node(src, SyntaxKind::END);
268 let range = environment_name_range(&end).expect("a name span");
269 assert_eq!(&src[range], "equation");
270 }
271
272 #[test]
273 fn environment_name_range_none_for_empty_name() {
274 assert_eq!(
275 environment_name_range(&node("\\begin{}\n\\end{}\n", SyntaxKind::BEGIN)),
276 None
277 );
278 }
279
280 #[test]
283 fn cast_is_kind_exact() {
284 let cmd = command("\\section{Hi}\n");
285 assert!(Command::cast(cmd.clone()).is_some());
286 assert!(Group::cast(cmd.clone()).is_none());
287 let group = nth_group(&cmd, 0).unwrap();
288 assert!(Group::cast(group.clone()).is_some());
289 assert!(Command::cast(group).is_none());
290 }
291
292 #[test]
293 fn typed_nth_group_is_a_group_node() {
294 let cmd = Command::cast(command("\\label{k}\n")).unwrap();
295 let group = cmd.nth_group(0).unwrap();
296 assert_eq!(group.syntax().kind(), SyntaxKind::GROUP);
297 }
298
299 #[test]
300 fn optionals_do_not_shift_group_indexing() {
301 let cmd = Command::cast(command("\\cmd[o]{a}\n")).unwrap();
303 assert_eq!(cmd.nth_group_text(0).as_deref(), Some("a"));
304 assert_eq!(cmd.optionals().count(), 1);
305 }
306
307 #[test]
308 fn first_group_range_stops_at_first_group() {
309 let src = "\\label{a}\n{b}\n";
311 let cmd = Command::cast(command(src)).unwrap();
312 assert_eq!(&src[cmd.first_group_range()], "\\label{a}");
313 assert_eq!(cmd.nth_group_text(1).as_deref(), Some("b"));
314 }
315
316 #[test]
317 fn free_fn_shims_stay_kind_agnostic() {
318 let begin = node(
323 "\\begin{macro}{\\foo}\ncode\n\\end{macro}\n",
324 SyntaxKind::BEGIN,
325 );
326 assert_eq!(
327 nth_group(&begin, 0).map(|g| g.kind()),
328 Some(SyntaxKind::GROUP)
329 );
330 assert_eq!(
331 group_command_name(&nth_group(&begin, 0).unwrap()).as_deref(),
332 Some("foo")
333 );
334 }
335
336 #[test]
337 fn environment_wrapper_reaches_begin_and_end() {
338 let env = Environment::cast(node(
339 "\\begin{equation}\nx\n\\end{equation}\n",
340 SyntaxKind::ENVIRONMENT,
341 ))
342 .unwrap();
343 assert_eq!(
344 env.begin().and_then(|b| b.name()).as_deref(),
345 Some("equation")
346 );
347 assert_eq!(
348 env.end().and_then(|e| e.name()).as_deref(),
349 Some("equation")
350 );
351 assert_eq!(env.name().as_deref(), Some("equation"));
352 }
353
354 #[test]
355 fn begin_name_falls_back_to_the_head_control_word() {
356 let src = "\\newcommand{\\bea}{\\begin{eqnarray}}\n\\newcommand{\\eea}{\\end{eqnarray}}\n\\bea a \\eea\n";
360 let begin = node(src, SyntaxKind::BEGIN);
361 assert_eq!(environment_name(&begin).as_deref(), Some("bea"));
362 assert_eq!(
363 Begin::cast(begin.clone()).unwrap().name().as_deref(),
364 Some("bea")
365 );
366 assert!(environment_name_range(&begin).is_none());
370 assert!(Begin::cast(begin).unwrap().name_range().is_none());
371 }
372
373 #[test]
374 fn alias_end_stays_nameless() {
375 let src = "\\newcommand{\\bea}{\\begin{eqnarray}}\n\\newcommand{\\eea}{\\end{eqnarray}}\n\\bea a \\eea\n";
378 let end = node(src, SyntaxKind::END);
379 assert!(environment_name(&end).is_none());
380 }
381
382 #[test]
383 fn a_real_begin_reads_its_name_group() {
384 let begin = node("\\begin{center}x\\end{center}", SyntaxKind::BEGIN);
388 assert_eq!(environment_name(&begin).as_deref(), Some("center"));
389 assert!(environment_name_range(&begin).is_some());
390 }
391}