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