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