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};
16use smol_str::SmolStr;
17
18use crate::syntax::{SyntaxKind, SyntaxNode, SyntaxToken};
19
20pub trait AstNode {
22 fn can_cast(kind: SyntaxKind) -> bool
23 where
24 Self: Sized;
25 fn cast(syntax: SyntaxNode) -> Option<Self>
26 where
27 Self: Sized;
28 fn syntax(&self) -> &SyntaxNode;
29}
30
31pub trait AstToken {
33 fn can_cast(kind: SyntaxKind) -> bool
34 where
35 Self: Sized;
36 fn cast(syntax: SyntaxToken) -> Option<Self>
37 where
38 Self: Sized;
39 fn syntax(&self) -> &SyntaxToken;
40 fn text(&self) -> &str {
41 self.syntax().text()
42 }
43}
44
45pub fn child<N: AstNode>(parent: &SyntaxNode) -> Option<N> {
47 parent.children().find_map(N::cast)
48}
49
50pub fn children<N: AstNode>(parent: &SyntaxNode) -> impl Iterator<Item = N> {
52 parent.children().filter_map(N::cast)
53}
54
55pub fn child_token<T: AstToken>(parent: &SyntaxNode) -> Option<T> {
57 parent
58 .children_with_tokens()
59 .filter_map(NodeOrToken::into_token)
60 .find_map(T::cast)
61}
62
63pub fn command_name(command: &SyntaxNode) -> Option<SmolStr> {
75 child_token::<ControlWord>(command).map(|cw| SmolStr::new(cw.name()))
76}
77
78pub fn control_word_range(command: &SyntaxNode) -> Option<TextRange> {
81 child_token::<ControlWord>(command).map(|cw| cw.range())
82}
83
84pub fn nth_group_text(command: &SyntaxNode, n: usize) -> Option<SmolStr> {
86 children::<Group>(command)
87 .nth(n)
88 .and_then(|g| g.inner_text())
89}
90
91pub fn nth_group_inner(command: &SyntaxNode, n: usize) -> Option<(TextRange, SmolStr)> {
94 children::<Group>(command).nth(n).and_then(|g| g.inner())
95}
96
97pub fn nth_group(command: &SyntaxNode, n: usize) -> Option<SyntaxNode> {
99 children::<Group>(command)
100 .nth(n)
101 .map(|g| g.syntax().clone())
102}
103
104pub fn first_group_range(command: &SyntaxNode) -> TextRange {
107 match children::<Group>(command).next() {
108 Some(group) => TextRange::new(
109 command.text_range().start(),
110 group.syntax().text_range().end(),
111 ),
112 None => command.text_range(),
113 }
114}
115
116pub fn group_command_name(group: &SyntaxNode) -> Option<SmolStr> {
121 if let Some(name) = child::<Command>(group).and_then(|c| c.name()) {
122 return Some(name);
123 }
124 let verb = group
125 .children_with_tokens()
126 .filter_map(|e| e.into_token())
127 .find(|t| t.kind() == SyntaxKind::VERB)?;
128 let name = verb.text().trim().strip_prefix('\\')?;
129 (!name.is_empty()
130 && name
131 .chars()
132 .all(|c| c.is_alphanumeric() || c == '@' || c == '_' || c == ':'))
133 .then(|| SmolStr::new(name))
134}
135
136pub fn group_inner_source(group: &SyntaxNode) -> String {
138 nodes::inner_source_of(group)
139}
140
141pub fn environment_name(begin_or_end: &SyntaxNode) -> Option<String> {
148 if begin_or_end.kind() == SyntaxKind::BEGIN {
149 return Begin::cast(begin_or_end.clone())?.name();
150 }
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";
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);
220 assert_eq!(nth_group_inner(&command("\\eqref{##1}\n"), 0), None);
221 assert_eq!(nth_group_text(&command("\\input{#1}\n"), 0), None);
222 }
223
224 #[test]
225 fn nth_group_text_none_for_nested_command() {
226 assert_eq!(nth_group_text(&command("\\input{\\jobname}\n"), 0), None);
227 }
228
229 #[test]
230 fn nth_group_text_none_when_group_absent() {
231 assert_eq!(nth_group_text(&command("\\input\n"), 0), None);
232 }
233
234 #[test]
235 fn group_command_name_reads_braced_control_word() {
236 let cmd = command("\\newcommand{\\foo}{x}\n");
237 let name = nth_group(&cmd, 0).and_then(|g| group_command_name(&g));
238 assert_eq!(name.as_deref(), Some("foo"));
239 }
240
241 #[test]
242 fn group_command_name_none_for_plain_text() {
243 let cmd = command("\\newenvironment{thm}{a}{b}\n");
244 let name = nth_group(&cmd, 0).and_then(|g| group_command_name(&g));
245 assert_eq!(name, None);
246 }
247
248 #[test]
249 fn group_inner_source_keeps_nested_braces() {
250 let cmd = command("\\NewDocumentCommand{\\foo}{m O{d} m}{x}\n");
251 let spec = nth_group(&cmd, 1).map(|g| group_inner_source(&g));
252 assert_eq!(spec.as_deref(), Some("m O{d} m"));
253 assert_eq!(nth_group_text(&cmd, 1), None);
254 }
255
256 #[test]
257 fn environment_name_range_spans_only_the_name() {
258 let src = "\\begin{equation}\nx\n\\end{equation}\n";
259 let begin = node(src, SyntaxKind::BEGIN);
260 let range = environment_name_range(&begin).expect("a name span");
261 assert_eq!(&src[range], "equation");
262
263 let end = node(src, SyntaxKind::END);
264 let range = environment_name_range(&end).expect("a name span");
265 assert_eq!(&src[range], "equation");
266 }
267
268 #[test]
269 fn environment_name_range_none_for_empty_name() {
270 assert_eq!(
271 environment_name_range(&node("\\begin{}\n\\end{}\n", SyntaxKind::BEGIN)),
272 None
273 );
274 }
275
276 #[test]
277 fn name_group_rejects_parameter_token() {
278 let mut builder = rowan::GreenNodeBuilder::new();
279 builder.start_node(SyntaxKind::NAME_GROUP.into());
280 builder.token(SyntaxKind::L_BRACE.into(), "{");
281 builder.token(SyntaxKind::HASH.into(), "#");
282 builder.token(SyntaxKind::WORD.into(), "1");
283 builder.token(SyntaxKind::R_BRACE.into(), "}");
284 builder.finish_node();
285 let group = NameGroup::cast(SyntaxNode::new_root(builder.finish())).unwrap();
286
287 assert_eq!(group.text(), None);
288 assert_eq!(group.range(), None);
289 }
290
291 #[test]
292 fn cast_is_kind_exact() {
293 let cmd = command("\\section{Hi}\n");
294 assert!(Command::cast(cmd.clone()).is_some());
295 assert!(Group::cast(cmd.clone()).is_none());
296 let group = nth_group(&cmd, 0).unwrap();
297 assert!(Group::cast(group.clone()).is_some());
298 assert!(Command::cast(group).is_none());
299 }
300
301 #[test]
302 fn typed_nth_group_is_a_group_node() {
303 let cmd = Command::cast(command("\\label{k}\n")).unwrap();
304 let group = cmd.nth_group(0).unwrap();
305 assert_eq!(group.syntax().kind(), SyntaxKind::GROUP);
306 }
307
308 #[test]
309 fn optionals_do_not_shift_group_indexing() {
310 let cmd = Command::cast(command("\\cmd[o]{a}\n")).unwrap();
311 assert_eq!(cmd.nth_group_text(0).as_deref(), Some("a"));
312 assert_eq!(cmd.optionals().count(), 1);
313 }
314
315 #[test]
316 fn first_group_range_stops_at_first_group() {
317 let src = "\\label{a}\n{b}\n";
318 let cmd = Command::cast(command(src)).unwrap();
319 assert_eq!(&src[cmd.first_group_range()], "\\label{a}");
320 assert_eq!(cmd.nth_group_text(1).as_deref(), Some("b"));
321 }
322
323 #[test]
324 fn free_fn_shims_stay_kind_agnostic() {
325 let begin = node(
326 "\\begin{macro}{\\foo}\ncode\n\\end{macro}\n",
327 SyntaxKind::BEGIN,
328 );
329 assert_eq!(
330 nth_group(&begin, 0).map(|g| g.kind()),
331 Some(SyntaxKind::GROUP)
332 );
333 assert_eq!(
334 group_command_name(&nth_group(&begin, 0).unwrap()).as_deref(),
335 Some("foo")
336 );
337 }
338
339 #[test]
340 fn environment_wrapper_reaches_begin_and_end() {
341 let env = Environment::cast(node(
342 "\\begin{equation}\nx\n\\end{equation}\n",
343 SyntaxKind::ENVIRONMENT,
344 ))
345 .unwrap();
346 assert_eq!(
347 env.begin().and_then(|b| b.name()).as_deref(),
348 Some("equation")
349 );
350 assert_eq!(
351 env.end().and_then(|e| e.name()).as_deref(),
352 Some("equation")
353 );
354 assert_eq!(env.name().as_deref(), Some("equation"));
355 }
356
357 #[test]
358 fn begin_name_falls_back_to_the_head_control_word() {
359 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());
367 assert!(Begin::cast(begin).unwrap().name_range().is_none());
368 }
369
370 #[test]
371 fn alias_end_stays_nameless() {
372 let src = "\\newcommand{\\bea}{\\begin{eqnarray}}\n\\newcommand{\\eea}{\\end{eqnarray}}\n\\bea a \\eea\n";
373 let end = node(src, SyntaxKind::END);
374 assert!(environment_name(&end).is_none());
375 }
376
377 #[test]
378 fn a_real_begin_reads_its_name_group() {
379 let begin = node("\\begin{center}x\\end{center}", SyntaxKind::BEGIN);
380 assert_eq!(environment_name(&begin).as_deref(), Some("center"));
381 assert!(environment_name_range(&begin).is_some());
382 }
383}