1use std::collections::BTreeSet;
10
11use thiserror::Error;
12
13use crate::token::{Token, TokenStoreError};
14use crate::{
15 CommonTokenStream, InputStream, Node, NodeId, NodeKind, Recognizer, TOKEN_EOF, Vocabulary,
16};
17
18mod generated {
19 pub(super) mod x_path_lexer;
20}
21
22use generated::x_path_lexer::{ANYWHERE, BANG, ID, ROOT, STRING, WILDCARD, XPathLexer};
23
24#[derive(Clone, Debug)]
26pub struct XPath {
27 path: String,
28 elements: Vec<PathElement>,
29}
30
31impl XPath {
32 pub fn new<R>(recognizer: &R, path: &str) -> Result<Self, XPathError>
34 where
35 R: Recognizer + ?Sized,
36 {
37 let tokens = tokenize(path)?;
38 let elements = compile_elements(&tokens, recognizer.rule_names(), recognizer.vocabulary())?;
39 if elements.is_empty() {
40 return Err(XPathError::MissingPathElement);
41 }
42 Ok(Self {
43 path: path.to_owned(),
44 elements,
45 })
46 }
47
48 #[must_use]
50 pub fn path(&self) -> &str {
51 &self.path
52 }
53
54 #[must_use]
61 pub fn evaluate<'tree>(&self, root: Node<'tree>) -> Vec<Node<'tree>> {
62 let mut work = vec![EvaluationNode::VirtualRoot(root)];
63 for element in &self.elements {
64 let mut next = Vec::new();
65 let mut seen = BTreeSet::new();
66 for node in work {
67 if !node.has_children() {
68 continue;
69 }
70 extend_unique(&mut next, &mut seen, evaluate_element(*element, node));
71 }
72 work = next;
73 }
74 work.into_iter()
75 .filter_map(EvaluationNode::tree_node)
76 .collect()
77 }
78
79 pub fn find_all<'tree, R>(
81 root: Node<'tree>,
82 path: &str,
83 recognizer: &R,
84 ) -> Result<Vec<Node<'tree>>, XPathError>
85 where
86 R: Recognizer + ?Sized,
87 {
88 Ok(Self::new(recognizer, path)?.evaluate(root))
89 }
90}
91
92#[derive(Clone, Debug, Eq, Error, PartialEq)]
94pub enum XPathError {
95 #[error("Invalid tokens or characters at index {index} in path '{path}'")]
96 InvalidCharacters { index: usize, path: String },
97 #[error("Missing path element at end of path")]
98 MissingPathElement,
99 #[error("{name} at index {index} isn't a valid token name")]
100 InvalidTokenName { name: String, index: usize },
101 #[error("{name} at index {index} isn't a valid rule name")]
102 InvalidRuleName { name: String, index: usize },
103 #[error("Unknown path element {element} at index {index}")]
104 UnknownPathElement { element: String, index: usize },
105 #[error("Could not tokenize path: {message}")]
106 Tokenization { message: String },
107}
108
109#[derive(Clone, Copy, Debug)]
110struct PathElement {
111 axis: Axis,
112 test: NodeTest,
113 invert: bool,
114}
115
116#[derive(Clone, Copy, Debug)]
117enum Axis {
118 Child,
119 DescendantOrSelf,
120}
121
122#[derive(Clone, Copy, Debug)]
123enum NodeTest {
124 Rule(usize),
125 Token(i32),
126 Wildcard,
127}
128
129#[derive(Clone, Copy, Debug)]
130enum EvaluationNode<'tree> {
131 VirtualRoot(Node<'tree>),
132 Tree(Node<'tree>),
133}
134
135impl<'tree> EvaluationNode<'tree> {
136 fn has_children(self) -> bool {
137 match self {
138 Self::VirtualRoot(_) => true,
139 Self::Tree(node) => node.children().next().is_some(),
140 }
141 }
142
143 const fn id(self) -> Option<NodeId> {
144 match self {
145 Self::VirtualRoot(_) => None,
146 Self::Tree(node) => Some(node.id()),
147 }
148 }
149
150 const fn tree_node(self) -> Option<Node<'tree>> {
151 match self {
152 Self::VirtualRoot(_) => None,
153 Self::Tree(node) => Some(node),
154 }
155 }
156}
157
158#[derive(Clone, Copy, Debug, Eq, PartialEq)]
159enum LexemeKind {
160 Anywhere,
161 Root,
162 Wildcard,
163 Bang,
164 Identifier,
165 String,
166 Eof,
167}
168
169#[derive(Clone, Debug)]
170struct Lexeme {
171 kind: LexemeKind,
172 text: String,
173 index: usize,
174}
175
176fn tokenize(path: &str) -> Result<Vec<Lexeme>, XPathError> {
177 let mut lexer = XPathLexer::new(InputStream::new(path));
178 lexer.remove_error_listeners();
179 let mut stream =
180 CommonTokenStream::try_new(lexer).map_err(|error| tokenization_error(&error))?;
181 stream.fill();
182
183 if let Some(error) = stream.drain_source_errors().into_iter().next() {
184 return Err(XPathError::InvalidCharacters {
185 index: source_index(path, error.line, error.column),
186 path: path.to_owned(),
187 });
188 }
189
190 stream
191 .tokens()
192 .map(|token| {
193 let kind = match token.token_type() {
194 ANYWHERE => LexemeKind::Anywhere,
195 ROOT => LexemeKind::Root,
196 WILDCARD => LexemeKind::Wildcard,
197 BANG => LexemeKind::Bang,
198 ID => LexemeKind::Identifier,
199 STRING => LexemeKind::String,
200 TOKEN_EOF => LexemeKind::Eof,
201 token_type => {
202 return Err(XPathError::Tokenization {
203 message: format!("unexpected XPath lexer token type {token_type}"),
204 });
205 }
206 };
207 Ok(Lexeme {
208 kind,
209 text: token.text_or_empty().to_owned(),
210 index: token.start(),
211 })
212 })
213 .collect()
214}
215
216fn tokenization_error(error: &TokenStoreError) -> XPathError {
217 XPathError::Tokenization {
218 message: error.to_string(),
219 }
220}
221
222fn source_index(path: &str, line: usize, column: usize) -> usize {
223 let mut current_line = 1;
224 let mut line_start = 0;
225 for (index, ch) in path.chars().enumerate() {
226 if current_line == line {
227 return line_start + column;
228 }
229 if ch == '\n' {
230 current_line += 1;
231 line_start = index + 1;
232 }
233 }
234 line_start + column
235}
236
237fn compile_elements(
238 tokens: &[Lexeme],
239 rule_names: &[String],
240 vocabulary: &Vocabulary,
241) -> Result<Vec<PathElement>, XPathError> {
242 let mut elements = Vec::new();
243 let mut cursor = 0;
244 while let Some(token) = tokens.get(cursor) {
245 match token.kind {
246 LexemeKind::Anywhere | LexemeKind::Root => {
247 let axis = if token.kind == LexemeKind::Anywhere {
248 Axis::DescendantOrSelf
249 } else {
250 Axis::Child
251 };
252 cursor += 1;
253 let invert = tokens
254 .get(cursor)
255 .is_some_and(|next| next.kind == LexemeKind::Bang);
256 cursor += usize::from(invert);
257 let word = tokens.get(cursor).ok_or(XPathError::MissingPathElement)?;
258 if word.kind == LexemeKind::Eof {
259 return Err(XPathError::MissingPathElement);
260 }
261 elements.push(compile_element(word, axis, invert, rule_names, vocabulary)?);
262 cursor += 1;
263 }
264 LexemeKind::Identifier | LexemeKind::Wildcard => {
265 elements.push(compile_element(
266 token,
267 Axis::Child,
268 false,
269 rule_names,
270 vocabulary,
271 )?);
272 cursor += 1;
273 }
274 LexemeKind::Eof => break,
275 LexemeKind::Bang | LexemeKind::String => {
276 return Err(XPathError::UnknownPathElement {
277 element: token.text.clone(),
278 index: token.index,
279 });
280 }
281 }
282 }
283 Ok(elements)
284}
285
286fn compile_element(
287 token: &Lexeme,
288 axis: Axis,
289 invert: bool,
290 rule_names: &[String],
291 vocabulary: &Vocabulary,
292) -> Result<PathElement, XPathError> {
293 let node_test = match token.kind {
294 LexemeKind::Wildcard => NodeTest::Wildcard,
295 LexemeKind::String => NodeTest::Token(resolve_token(token, vocabulary)?),
296 LexemeKind::Identifier if token.text.starts_with(char::is_uppercase) => {
297 NodeTest::Token(resolve_token(token, vocabulary)?)
298 }
299 _ => NodeTest::Rule(resolve_rule(token, rule_names)?),
300 };
301 Ok(PathElement {
302 axis,
303 test: node_test,
304 invert,
305 })
306}
307
308fn resolve_token(token: &Lexeme, vocabulary: &Vocabulary) -> Result<i32, XPathError> {
309 vocabulary
310 .token_type(&token.text)
311 .ok_or_else(|| XPathError::InvalidTokenName {
312 name: token.text.clone(),
313 index: token.index,
314 })
315}
316
317fn resolve_rule(token: &Lexeme, rule_names: &[String]) -> Result<usize, XPathError> {
318 rule_names
319 .iter()
320 .rposition(|name| name == &token.text)
321 .ok_or_else(|| XPathError::InvalidRuleName {
322 name: token.text.clone(),
323 index: token.index,
324 })
325}
326
327fn evaluate_element<'tree>(
328 element: PathElement,
329 node: EvaluationNode<'tree>,
330) -> Box<dyn Iterator<Item = EvaluationNode<'tree>> + 'tree> {
331 match node {
332 EvaluationNode::VirtualRoot(root) => evaluate_virtual_root(element, root),
333 EvaluationNode::Tree(node) => evaluate_tree_element(element, node),
334 }
335}
336
337fn evaluate_virtual_root<'tree>(
338 element: PathElement,
339 root: Node<'tree>,
340) -> Box<dyn Iterator<Item = EvaluationNode<'tree>> + 'tree> {
341 match element.axis {
342 Axis::Child => Box::new(
343 std::iter::once(root)
344 .filter(move |node| matches_direct(*node, element.test, element.invert))
345 .map(EvaluationNode::Tree),
346 ),
347 Axis::DescendantOrSelf if matches!(element.test, NodeTest::Wildcard) && !element.invert => {
348 Box::new(
349 std::iter::once(EvaluationNode::VirtualRoot(root))
350 .chain(evaluate_anywhere(element, root).map(EvaluationNode::Tree)),
351 )
352 }
353 Axis::DescendantOrSelf => {
354 Box::new(evaluate_anywhere(element, root).map(EvaluationNode::Tree))
355 }
356 }
357}
358
359fn evaluate_tree_element<'tree>(
360 element: PathElement,
361 node: Node<'tree>,
362) -> Box<dyn Iterator<Item = EvaluationNode<'tree>> + 'tree> {
363 match element.axis {
364 Axis::Child => Box::new(
365 node.children()
366 .filter(move |child| matches_direct(*child, element.test, element.invert))
367 .map(EvaluationNode::Tree),
368 ),
369 Axis::DescendantOrSelf => {
370 Box::new(evaluate_anywhere(element, node).map(EvaluationNode::Tree))
371 }
372 }
373}
374
375fn evaluate_anywhere<'tree>(
376 element: PathElement,
377 node: Node<'tree>,
378) -> Box<dyn Iterator<Item = Node<'tree>> + 'tree> {
379 match element.test {
380 NodeTest::Wildcard if element.invert => Box::new(std::iter::empty()),
381 NodeTest::Wildcard => Box::new(node.descendants()),
382 NodeTest::Rule(rule_index) => Box::new(node.descendants().filter(move |candidate| {
383 candidate
384 .as_rule()
385 .is_some_and(|rule| rule.rule_index() == rule_index)
386 })),
387 NodeTest::Token(token_type) => Box::new(
388 node.descendants()
389 .filter(move |candidate| node_token_type(*candidate) == Some(token_type)),
390 ),
391 }
392}
393
394fn matches_direct(node: Node<'_>, node_test: NodeTest, invert: bool) -> bool {
395 match node_test {
396 NodeTest::Wildcard => !invert,
397 NodeTest::Rule(rule_index) => node
398 .as_rule()
399 .is_some_and(|rule| (rule.rule_index() == rule_index) != invert),
400 NodeTest::Token(token_type) => {
401 node_token_type(node).is_some_and(|actual| (actual == token_type) != invert)
402 }
403 }
404}
405
406fn node_token_type(node: Node<'_>) -> Option<i32> {
407 match node.kind() {
408 NodeKind::Terminal => node
409 .as_terminal()
410 .map(|terminal| terminal.symbol().token_type()),
411 NodeKind::Error => node.as_error().map(|error| error.symbol().token_type()),
412 NodeKind::Rule => None,
413 }
414}
415
416fn extend_unique<'tree>(
417 nodes: &mut Vec<EvaluationNode<'tree>>,
418 seen: &mut BTreeSet<Option<NodeId>>,
419 matches: impl IntoIterator<Item = EvaluationNode<'tree>>,
420) {
421 for node in matches {
422 if seen.insert(node.id()) {
423 nodes.push(node);
424 }
425 }
426}
427
428#[cfg(test)]
429#[allow(clippy::disallowed_methods)] mod tests {
431 use super::*;
432 use crate::RecognizerData;
433 use crate::token::{TokenSpec, TokenStore};
434 use crate::tree::{ParseTreeStorage, ParsedFile, ParserRuleContext};
435
436 const PROG: usize = 0;
437 const FUNC: usize = 1;
438 const BODY: usize = 2;
439 const ARG: usize = 3;
440 const STAT: usize = 4;
441 const EXPR: usize = 5;
442 const PRIMARY: usize = 6;
443
444 const DEF: i32 = 1;
445 const LPAREN: i32 = 2;
446 const COMMA: i32 = 3;
447 const RPAREN: i32 = 4;
448 const LBRACE: i32 = 5;
449 const RBRACE: i32 = 6;
450 const SEMI: i32 = 7;
451 const ASSIGN: i32 = 8;
452 const MUL: i32 = 9;
453 const ADD: i32 = 11;
454 const RETURN: i32 = 13;
455 const IDENTIFIER: i32 = 14;
456 const INTEGER: i32 = 15;
457
458 enum TreeSpec {
459 Rule(usize, Vec<Self>),
460 Token(i32, &'static str),
461 }
462
463 fn rule(index: usize, children: Vec<TreeSpec>) -> TreeSpec {
464 TreeSpec::Rule(index, children)
465 }
466
467 const fn token(token_type: i32, text: &'static str) -> TreeSpec {
468 TreeSpec::Token(token_type, text)
469 }
470
471 fn primary_expr(token_type: i32, text: &'static str) -> TreeSpec {
472 rule(EXPR, vec![rule(PRIMARY, vec![token(token_type, text)])])
473 }
474
475 fn binary_expr(left: TreeSpec, operator: TreeSpec, right: TreeSpec) -> TreeSpec {
476 rule(EXPR, vec![left, operator, right])
477 }
478
479 fn first_function() -> TreeSpec {
480 let assignment = rule(
481 STAT,
482 vec![
483 token(IDENTIFIER, "x"),
484 token(ASSIGN, "="),
485 binary_expr(
486 primary_expr(INTEGER, "3"),
487 token(ADD, "+"),
488 primary_expr(INTEGER, "4"),
489 ),
490 token(SEMI, ";"),
491 ],
492 );
493 let print = rule(STAT, vec![primary_expr(IDENTIFIER, "y"), token(SEMI, ";")]);
494 let body = rule(
495 BODY,
496 vec![
497 token(LBRACE, "{"),
498 assignment,
499 print,
500 rule(STAT, vec![token(SEMI, ";")]),
501 token(RBRACE, "}"),
502 ],
503 );
504 rule(
505 FUNC,
506 vec![
507 token(DEF, "def"),
508 token(IDENTIFIER, "f"),
509 token(LPAREN, "("),
510 rule(ARG, vec![token(IDENTIFIER, "x")]),
511 token(COMMA, ","),
512 rule(ARG, vec![token(IDENTIFIER, "y")]),
513 token(RPAREN, ")"),
514 body,
515 ],
516 )
517 }
518
519 fn second_function() -> TreeSpec {
520 let product = binary_expr(
521 primary_expr(INTEGER, "2"),
522 token(MUL, "*"),
523 primary_expr(IDENTIFIER, "x"),
524 );
525 let returned = binary_expr(primary_expr(INTEGER, "1"), token(ADD, "+"), product);
526 let body = rule(
527 BODY,
528 vec![
529 token(LBRACE, "{"),
530 rule(
531 STAT,
532 vec![token(RETURN, "return"), returned, token(SEMI, ";")],
533 ),
534 token(RBRACE, "}"),
535 ],
536 );
537 rule(
538 FUNC,
539 vec![
540 token(DEF, "def"),
541 token(IDENTIFIER, "g"),
542 token(LPAREN, "("),
543 rule(ARG, vec![token(IDENTIFIER, "x")]),
544 token(RPAREN, ")"),
545 body,
546 ],
547 )
548 }
549
550 fn materialize(
551 spec: TreeSpec,
552 tokens: &mut TokenStore,
553 storage: &mut ParseTreeStorage,
554 ) -> NodeId {
555 match spec {
556 TreeSpec::Token(token_type, text) => {
557 let token = tokens
558 .push(TokenSpec::explicit(token_type, text))
559 .expect("test token should fit");
560 storage.terminal(token)
561 }
562 TreeSpec::Rule(rule_index, children) => {
563 let mut context = ParserRuleContext::new(rule_index, -1);
564 for child in children {
565 let child = materialize(child, tokens, storage);
566 storage.add_child(&mut context, child);
567 }
568 storage.finish_rule(context)
569 }
570 }
571 }
572
573 fn sample_tree() -> ParsedFile {
574 let mut tokens = TokenStore::new(None, "Expr");
575 let mut storage = ParseTreeStorage::new();
576 let root = materialize(
577 rule(PROG, vec![first_function(), second_function()]),
578 &mut tokens,
579 &mut storage,
580 );
581 ParsedFile::new(tokens, storage, root)
582 }
583
584 #[derive(Debug)]
585 struct TestRecognizer {
586 data: RecognizerData,
587 }
588
589 impl Recognizer for TestRecognizer {
590 fn data(&self) -> &RecognizerData {
591 &self.data
592 }
593
594 fn data_mut(&mut self) -> &mut RecognizerData {
595 &mut self.data
596 }
597 }
598
599 fn expr_recognizer() -> TestRecognizer {
600 let vocabulary = Vocabulary::new(
601 [
602 None,
603 Some("'def'"),
604 Some("'('"),
605 Some("','"),
606 Some("')'"),
607 Some("'{'"),
608 Some("'}'"),
609 Some("';'"),
610 Some("'='"),
611 Some("'*'"),
612 Some("'/'"),
613 Some("'+'"),
614 Some("'-'"),
615 Some("'return'"),
616 None,
617 None,
618 ],
619 [
620 None,
621 None,
622 None,
623 None,
624 None,
625 None,
626 None,
627 None,
628 None,
629 Some("MUL"),
630 Some("DIV"),
631 Some("ADD"),
632 Some("SUB"),
633 Some("RETURN"),
634 Some("ID"),
635 Some("INT"),
636 ],
637 [None::<&str>; 16],
638 );
639 TestRecognizer {
640 data: RecognizerData::new("Expr.g4", vocabulary)
641 .with_rule_names(["prog", "func", "body", "arg", "stat", "expr", "primary"]),
642 }
643 }
644
645 fn display_nodes(nodes: Vec<Node<'_>>, rule_names: &[String]) -> Vec<String> {
646 nodes
647 .into_iter()
648 .map(|node| {
649 node.as_rule()
650 .map_or_else(|| node.text(), |rule| rule_names[rule.rule_index()].clone())
651 })
652 .collect()
653 }
654
655 #[test]
656 fn upstream_valid_paths() {
657 let tree = sample_tree();
658 let recognizer = expr_recognizer();
659 let paths = [
660 "/prog/func",
661 "/prog/*",
662 "/*/func",
663 "prog",
664 "/prog",
665 "/*",
666 "*",
667 "//ID",
668 "//expr/primary/ID",
669 "//body//ID",
670 "//'return'",
671 "//RETURN",
672 "//primary/*",
673 "//func/*/stat",
674 "/prog/func/'def'",
675 "//stat/';'",
676 "//expr/primary/!ID",
677 "//expr/!primary",
678 "//!*",
679 "/!*",
680 "//expr//ID",
681 ];
682 let results = paths
683 .into_iter()
684 .map(|path| {
685 let nodes =
686 XPath::find_all(tree.tree(), path, &recognizer).expect("valid XPath query");
687 (path, display_nodes(nodes, recognizer.data().rule_names()))
688 })
689 .collect::<Vec<_>>();
690
691 insta::assert_debug_snapshot!("upstream_valid_paths", results);
692 }
693
694 #[test]
695 fn upstream_invalid_paths() {
696 let recognizer = expr_recognizer();
697 let paths = ["&", "//w&e/", "///", "//", "//Ick", "/prog/ick"];
698 let errors = paths
699 .into_iter()
700 .map(|path| {
701 (
702 path,
703 XPath::new(&recognizer, path)
704 .expect_err("invalid XPath query")
705 .to_string(),
706 )
707 })
708 .collect::<Vec<_>>();
709
710 insta::assert_debug_snapshot!("upstream_invalid_paths", errors);
711 }
712
713 #[test]
714 fn lexer_and_parser_edge_cases_are_explicit() {
715 let recognizer = expr_recognizer();
716 let paths = ["", "// ID", "//'return", "!", "//!"];
717 let errors = paths
718 .into_iter()
719 .map(|path| {
720 (
721 path,
722 XPath::new(&recognizer, path)
723 .expect_err("invalid XPath query")
724 .to_string(),
725 )
726 })
727 .collect::<Vec<_>>();
728
729 insta::assert_debug_snapshot!("lexer_and_parser_edge_cases", errors);
730 }
731
732 #[test]
733 fn generated_lexer_accepts_upstream_unicode_name_ranges() {
734 let recognizer = TestRecognizer {
735 data: RecognizerData::new(
736 "Unicode.g4",
737 Vocabulary::new([None::<&str>; 2], [None, Some("ÄTOKEN")], [None::<&str>; 2]),
738 )
739 .with_rule_names(["éclair", "文"]),
740 };
741
742 assert!(XPath::new(&recognizer, "//ÄTOKEN").is_ok());
743 assert!(XPath::new(&recognizer, "//éclair").is_ok());
744 assert!(XPath::new(&recognizer, "//文").is_ok());
745 }
746
747 #[test]
748 fn named_anywhere_inversion_matches_java_4_13_2() {
749 let tree = sample_tree();
750 let recognizer = expr_recognizer();
751 let paths = ["//ID", "//!ID", "//expr", "//!expr"];
752 let results = paths
753 .into_iter()
754 .map(|path| {
755 let nodes =
756 XPath::find_all(tree.tree(), path, &recognizer).expect("valid XPath query");
757 (path, display_nodes(nodes, recognizer.data().rule_names()))
758 })
759 .collect::<Vec<_>>();
760
761 insta::assert_debug_snapshot!("named_anywhere_inversion_matches_java_4_13_2", results);
762 }
763
764 #[test]
765 fn wildcard_anywhere_preserves_java_virtual_root_for_later_steps() {
766 let tree = sample_tree();
767 let recognizer = expr_recognizer();
768 let nodes =
769 XPath::find_all(tree.tree(), "//*/prog", &recognizer).expect("valid XPath query");
770
771 insta::assert_debug_snapshot!(
772 "wildcard_anywhere_virtual_root",
773 display_nodes(nodes, recognizer.data().rule_names())
774 );
775 }
776}