1use crate::errors::AntlrError;
2use crate::recognizer::Recognizer;
3use std::any::Any;
4use std::fmt;
5use std::rc::Rc;
6
7use crate::token::{CommonToken, Token, TokenRef};
8use std::collections::BTreeMap;
9
10#[derive(Clone, Debug, Eq, PartialEq)]
11pub enum ParseTree {
12 Rule(RuleNode),
13 Terminal(TerminalNode),
14 Error(ErrorNode),
15}
16
17impl ParseTree {
18 pub fn text(&self) -> String {
19 match self {
20 Self::Rule(rule) => rule.text(),
21 Self::Terminal(node) => node.text(),
22 Self::Error(node) => node.text(),
23 }
24 }
25
26 pub fn to_string_tree_with_names<S: AsRef<str>>(&self, rule_names: &[S]) -> String {
27 match self {
28 Self::Rule(rule) => rule.to_string_tree_with_names(rule_names),
29 Self::Terminal(node) => escape_tree_text(&node.text()),
30 Self::Error(node) => escape_tree_text(&node.text()),
31 }
32 }
33
34 pub fn to_string_tree<R: Recognizer>(&self, recognizer: Option<&R>) -> String {
38 recognizer.map_or_else(
39 || self.to_string_tree_with_names::<&str>(&[]),
40 |recognizer| self.to_string_tree_with_names(recognizer.data().rule_names()),
41 )
42 }
43
44 pub fn first_rule(&self, rule_index: usize) -> Option<&Self> {
46 match self {
47 Self::Rule(rule) => {
48 if rule.context().rule_index() == rule_index {
49 return Some(self);
50 }
51 rule.context()
52 .children()
53 .iter()
54 .find_map(|child| child.first_rule(rule_index))
55 }
56 Self::Terminal(_) | Self::Error(_) => None,
57 }
58 }
59
60 pub fn first_rule_stop(&self, rule_index: usize) -> Option<&CommonToken> {
62 let Self::Rule(rule) = self else {
63 return None;
64 };
65 if rule.context().rule_index() == rule_index {
66 return rule.context().stop();
67 }
68 rule.context()
69 .children()
70 .iter()
71 .find_map(|child| child.first_rule_stop(rule_index))
72 }
73
74 pub fn first_rule_int_return(&self, rule_index: usize, name: &str) -> Option<i64> {
78 let Self::Rule(rule) = self else {
79 return None;
80 };
81 if rule.context().rule_index() == rule_index {
82 return rule.context().int_return(name);
83 }
84 rule.context()
85 .children()
86 .iter()
87 .find_map(|child| child.first_rule_int_return(rule_index, name))
88 }
89
90 pub fn rule_attrs<T: Any>(&self) -> Option<&T> {
95 let Self::Rule(rule) = self else {
96 return None;
97 };
98 rule.context().generated_attrs::<T>()
99 }
100
101 pub fn first_error_token(&self) -> Option<&CommonToken> {
103 match self {
104 Self::Rule(rule) => rule
105 .context()
106 .children()
107 .iter()
108 .find_map(Self::first_error_token),
109 Self::Terminal(_) => None,
110 Self::Error(node) => Some(node.symbol()),
111 }
112 }
113
114 pub fn rule_invocation_stack<S: AsRef<str>>(
117 &self,
118 rule_index: usize,
119 rule_names: &[S],
120 ) -> Option<Vec<String>> {
121 let mut stack = Vec::new();
122 if self.find_rule_path(rule_index, rule_names, &mut stack) {
123 stack.reverse();
124 return Some(stack);
125 }
126 None
127 }
128
129 fn find_rule_path<S: AsRef<str>>(
130 &self,
131 rule_index: usize,
132 rule_names: &[S],
133 stack: &mut Vec<String>,
134 ) -> bool {
135 let Self::Rule(rule) = self else {
136 return false;
137 };
138 let current_index = rule.context().rule_index();
139 stack.push(
140 rule_names
141 .get(current_index)
142 .map_or("<unknown>", |name| name.as_ref())
143 .to_owned(),
144 );
145 if current_index == rule_index
146 || rule
147 .context()
148 .children()
149 .iter()
150 .any(|child| child.find_rule_path(rule_index, rule_names, stack))
151 {
152 return true;
153 }
154 stack.pop();
155 false
156 }
157}
158
159fn escape_tree_text(text: &str) -> String {
160 let mut escaped = String::with_capacity(text.len());
161 for ch in text.chars() {
162 match ch {
163 '\n' => escaped.push_str("\\n"),
164 '\r' => escaped.push_str("\\r"),
165 '\t' => escaped.push_str("\\t"),
166 _ => escaped.push(ch),
167 }
168 }
169 escaped
170}
171
172#[derive(Clone, Debug, Eq, PartialEq)]
173pub struct RuleNode {
174 context: ParserRuleContext,
175}
176
177impl RuleNode {
178 pub const fn new(context: ParserRuleContext) -> Self {
179 Self { context }
180 }
181
182 pub const fn context(&self) -> &ParserRuleContext {
183 &self.context
184 }
185
186 pub const fn context_mut(&mut self) -> &mut ParserRuleContext {
187 &mut self.context
188 }
189
190 pub fn text(&self) -> String {
191 self.context.text()
192 }
193
194 pub fn to_string_tree_with_names<S: AsRef<str>>(&self, rule_names: &[S]) -> String {
195 self.context.to_string_tree_with_names(rule_names)
196 }
197}
198
199#[derive(Clone, Debug, Eq, PartialEq)]
200pub struct ParserRuleContext {
201 rule_index: usize,
202 invoking_state: isize,
203 alt_number: usize,
204 start: Option<TokenRef>,
205 stop: Option<TokenRef>,
206 int_returns: Option<Box<IntReturns>>,
207 children: Vec<ParseTree>,
208 matched_child: bool,
214 exception: Option<Box<AntlrError>>,
218 attrs: Option<GeneratedAttrs>,
220}
221
222#[derive(Clone, Debug, Default, Eq, PartialEq)]
223struct IntReturns(BTreeMap<String, i64>);
224
225#[derive(Clone)]
233pub struct GeneratedAttrs(Rc<dyn Any>);
234
235impl GeneratedAttrs {
236 pub fn new<T: Any>(attrs: T) -> Self {
237 Self(Rc::new(attrs))
238 }
239
240 pub fn downcast_ref<T: Any>(&self) -> Option<&T> {
241 self.0.downcast_ref::<T>()
242 }
243}
244
245impl fmt::Debug for GeneratedAttrs {
246 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
247 f.write_str("GeneratedAttrs(..)")
248 }
249}
250
251impl PartialEq for GeneratedAttrs {
254 fn eq(&self, other: &Self) -> bool {
255 Rc::ptr_eq(&self.0, &other.0)
256 }
257}
258
259impl Eq for GeneratedAttrs {}
260
261impl ParserRuleContext {
262 pub const fn new(rule_index: usize, invoking_state: isize) -> Self {
263 Self {
264 rule_index,
265 invoking_state,
266 alt_number: 0,
267 start: None,
268 stop: None,
269 int_returns: None,
270 children: Vec::new(),
271 matched_child: false,
272 exception: None,
273 attrs: None,
274 }
275 }
276
277 pub(crate) fn with_child_capacity(
278 rule_index: usize,
279 invoking_state: isize,
280 capacity: usize,
281 ) -> Self {
282 Self {
283 rule_index,
284 invoking_state,
285 alt_number: 0,
286 start: None,
287 stop: None,
288 int_returns: None,
289 children: Vec::with_capacity(capacity),
290 matched_child: false,
291 exception: None,
292 attrs: None,
293 }
294 }
295
296 pub const fn rule_index(&self) -> usize {
297 self.rule_index
298 }
299
300 pub const fn invoking_state(&self) -> isize {
301 self.invoking_state
302 }
303
304 pub const fn alt_number(&self) -> usize {
305 self.alt_number
306 }
307
308 pub const fn set_alt_number(&mut self, alt_number: usize) {
309 self.alt_number = alt_number;
310 }
311
312 pub fn start(&self) -> Option<&CommonToken> {
313 self.start.as_deref()
314 }
315
316 pub(crate) fn start_ref(&self) -> Option<TokenRef> {
317 self.start.as_ref().map(Rc::clone)
318 }
319
320 pub fn stop(&self) -> Option<&CommonToken> {
321 self.stop.as_deref()
322 }
323
324 pub fn set_start(&mut self, token: CommonToken) {
325 self.start = Some(Rc::new(token));
326 }
327
328 pub(crate) fn set_start_ref(&mut self, token: TokenRef) {
329 self.start = Some(token);
330 }
331
332 pub fn set_stop(&mut self, token: CommonToken) {
333 self.stop = Some(Rc::new(token));
334 }
335
336 pub(crate) fn set_stop_ref(&mut self, token: TokenRef) {
337 self.stop = Some(token);
338 }
339
340 pub fn set_int_return(&mut self, name: impl Into<String>, value: i64) {
342 self.int_returns
343 .get_or_insert_with(Box::default)
344 .0
345 .insert(name.into(), value);
346 }
347
348 pub fn int_return(&self, name: &str) -> Option<i64> {
350 self.int_returns
351 .as_ref()
352 .and_then(|values| values.0.get(name).copied())
353 }
354
355 pub fn set_generated_attrs(&mut self, attrs: GeneratedAttrs) {
357 self.attrs = Some(attrs);
358 }
359
360 pub fn generated_attrs<T: Any>(&self) -> Option<&T> {
362 self.attrs.as_ref().and_then(GeneratedAttrs::downcast_ref)
363 }
364
365 pub fn exception(&self) -> Option<&AntlrError> {
366 self.exception.as_deref()
367 }
368
369 pub fn set_exception(&mut self, error: AntlrError) {
370 self.exception = Some(Box::new(error));
371 }
372
373 pub fn children(&self) -> &[ParseTree] {
374 &self.children
375 }
376
377 pub const fn child_count(&self) -> usize {
379 self.children.len()
380 }
381
382 pub fn child_rule(&self, rule_index: usize) -> Option<&Self> {
384 self.child_rules(rule_index).next()
385 }
386
387 pub fn child_rules(&self, rule_index: usize) -> impl Iterator<Item = &Self> + '_ {
389 self.children.iter().filter_map(move |child| match child {
390 ParseTree::Rule(rule) if rule.context().rule_index() == rule_index => {
391 Some(rule.context())
392 }
393 _ => None,
394 })
395 }
396
397 pub fn child_token(&self, token_type: i32) -> Option<&TerminalNode> {
402 self.child_tokens(token_type).next()
403 }
404
405 pub fn child_rule_trees(&self, rule_index: usize) -> impl Iterator<Item = &ParseTree> + '_ {
411 self.children.iter().filter(move |child| match child {
412 ParseTree::Rule(rule) => rule.context().rule_index() == rule_index,
413 ParseTree::Terminal(_) | ParseTree::Error(_) => false,
414 })
415 }
416
417 pub fn child_tokens(&self, token_type: i32) -> impl Iterator<Item = &TerminalNode> + '_ {
420 self.children.iter().filter_map(move |child| match child {
421 ParseTree::Terminal(node) if node.symbol().token_type() == token_type => Some(node),
422 ParseTree::Error(node) if node.symbol().token_type() == token_type => {
423 Some(node.terminal())
424 }
425 _ => None,
426 })
427 }
428
429 pub fn terminal_children(&self) -> impl Iterator<Item = &TerminalNode> + '_ {
431 self.children.iter().filter_map(|child| match child {
432 ParseTree::Terminal(node) => Some(node),
433 ParseTree::Error(node) => Some(node.terminal()),
434 ParseTree::Rule(_) => None,
435 })
436 }
437
438 pub fn downcast_ref<T: FromRuleContext>(&self) -> Option<T> {
444 T::from_rule_context(self)
445 }
446
447 pub fn has_token(&self, token_type: i32) -> bool {
449 self.child_token(token_type).is_some()
450 }
451
452 pub fn add_child(&mut self, child: ParseTree) {
453 self.matched_child = true;
454 self.children.push(child);
455 }
456
457 pub const fn note_matched_child(&mut self) {
461 self.matched_child = true;
462 }
463
464 pub const fn has_matched_child(&self) -> bool {
467 self.matched_child
468 }
469
470 pub fn text(&self) -> String {
471 self.children.iter().map(ParseTree::text).collect()
472 }
473
474 pub fn to_string_tree_with_names<S: AsRef<str>>(&self, rule_names: &[S]) -> String {
475 let name = rule_names
476 .get(self.rule_index)
477 .map_or("<unknown>", |name| name.as_ref());
478 let display_name = if self.alt_number == 0 {
479 name.to_owned()
480 } else {
481 format!("{name}:{}", self.alt_number)
482 };
483 if self.children.is_empty() {
484 return display_name;
485 }
486 let children = self
487 .children
488 .iter()
489 .map(|child| child.to_string_tree_with_names(rule_names))
490 .collect::<Vec<_>>()
491 .join(" ");
492 format!("({display_name} {children})")
493 }
494
495 pub fn to_string_tree<R: Recognizer>(&self, recognizer: Option<&R>) -> String {
499 recognizer.map_or_else(
500 || self.to_string_tree_with_names::<&str>(&[]),
501 |recognizer| self.to_string_tree_with_names(recognizer.data().rule_names()),
502 )
503 }
504}
505
506#[derive(Clone, Debug, Eq, PartialEq)]
507pub struct TerminalNode {
508 token: TokenRef,
509}
510
511impl TerminalNode {
512 pub fn new(token: CommonToken) -> Self {
513 Self {
514 token: Rc::new(token),
515 }
516 }
517
518 pub(crate) const fn from_ref(token: TokenRef) -> Self {
519 Self { token }
520 }
521
522 pub fn symbol(&self) -> &CommonToken {
523 self.token.as_ref()
524 }
525
526 pub fn text(&self) -> String {
527 self.token.text().to_owned()
528 }
529}
530
531impl fmt::Display for TerminalNode {
534 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
535 f.write_str(self.token.text())
536 }
537}
538
539#[derive(Clone, Debug, Eq, PartialEq)]
540pub struct ErrorNode {
541 terminal: TerminalNode,
542}
543
544impl ErrorNode {
545 pub fn new(token: CommonToken) -> Self {
546 Self {
547 terminal: TerminalNode::new(token),
548 }
549 }
550
551 pub(crate) const fn from_ref(token: TokenRef) -> Self {
552 Self {
553 terminal: TerminalNode::from_ref(token),
554 }
555 }
556
557 const fn terminal(&self) -> &TerminalNode {
558 &self.terminal
559 }
560
561 pub fn symbol(&self) -> &CommonToken {
562 self.terminal.symbol()
563 }
564
565 pub fn text(&self) -> String {
566 self.terminal.text()
567 }
568}
569
570pub trait FromRuleContext: Sized {
577 fn from_rule_context(context: &ParserRuleContext) -> Option<Self>;
578}
579
580pub trait ParseTreeListener {
581 fn enter_every_rule(&mut self, _ctx: &ParserRuleContext) -> Result<(), AntlrError> {
582 Ok(())
583 }
584
585 fn exit_every_rule(&mut self, _ctx: &ParserRuleContext) -> Result<(), AntlrError> {
586 Ok(())
587 }
588
589 fn visit_terminal(&mut self, _node: &TerminalNode) -> Result<(), AntlrError> {
590 Ok(())
591 }
592
593 fn visit_error_node(&mut self, _node: &ErrorNode) -> Result<(), AntlrError> {
594 Ok(())
595 }
596}
597
598#[derive(Debug, Default)]
599pub struct ParseTreeWalker;
600
601impl ParseTreeWalker {
602 pub fn walk<L: ParseTreeListener>(
605 listener: &mut L,
606 tree: &ParseTree,
607 ) -> Result<(), AntlrError> {
608 match tree {
609 ParseTree::Rule(rule) => {
610 listener.enter_every_rule(rule.context())?;
611 for child in rule.context().children() {
612 Self::walk(listener, child)?;
613 }
614 listener.exit_every_rule(rule.context())
615 }
616 ParseTree::Terminal(node) => listener.visit_terminal(node),
617 ParseTree::Error(node) => listener.visit_error_node(node),
618 }
619 }
620}
621
622#[cfg(test)]
623mod tests {
624 use super::*;
625 use crate::token::CommonToken;
626
627 #[test]
628 fn renders_rule_tree() {
629 let mut ctx = ParserRuleContext::new(0, -1);
630 ctx.add_child(ParseTree::Terminal(TerminalNode::new(
631 CommonToken::new(1).with_text("x"),
632 )));
633 let tree = ParseTree::Rule(RuleNode::new(ctx));
634 assert_eq!(tree.to_string_tree_with_names(&["expr"]), "(expr x)");
635 }
636
637 #[test]
638 fn finds_first_rule_depth_first() {
639 let mut nested = ParserRuleContext::new(1, -1);
640 nested.add_child(ParseTree::Terminal(TerminalNode::new(
641 CommonToken::new(1).with_text("x"),
642 )));
643
644 let mut root = ParserRuleContext::new(0, -1);
645 root.add_child(ParseTree::Rule(RuleNode::new(nested)));
646 let tree = ParseTree::Rule(RuleNode::new(root));
647
648 let rule = tree.first_rule(1).expect("nested rule should be found");
649 assert_eq!(
650 rule.to_string_tree_with_names(&["root".to_owned(), "child".to_owned()]),
651 "(child x)"
652 );
653 assert!(tree.first_rule(2).is_none());
654 }
655
656 #[test]
657 fn reports_rule_invocation_stack_from_leaf_to_root() {
658 let mut nested = ParserRuleContext::new(1, -1);
659 nested.add_child(ParseTree::Terminal(TerminalNode::new(
660 CommonToken::new(1).with_text("x"),
661 )));
662
663 let mut root = ParserRuleContext::new(0, -1);
664 root.add_child(ParseTree::Rule(RuleNode::new(nested)));
665 let tree = ParseTree::Rule(RuleNode::new(root));
666
667 assert_eq!(
668 tree.rule_invocation_stack(1, &["s", "a"]),
669 Some(vec!["a".to_owned(), "s".to_owned()])
670 );
671 }
672
673 #[test]
674 fn finds_direct_child_rules_by_index() {
675 let mut direct = ParserRuleContext::new(1, -1);
676 direct.add_child(ParseTree::Terminal(TerminalNode::new(
677 CommonToken::new(10).with_text("direct"),
678 )));
679
680 let mut nested_match = ParserRuleContext::new(1, -1);
681 nested_match.add_child(ParseTree::Terminal(TerminalNode::new(
682 CommonToken::new(11).with_text("nested"),
683 )));
684 let mut wrapper = ParserRuleContext::new(2, -1);
685 wrapper.add_child(ParseTree::Rule(RuleNode::new(nested_match)));
686
687 let mut root = ParserRuleContext::new(0, -1);
688 root.add_child(ParseTree::Terminal(TerminalNode::new(
689 CommonToken::new(12).with_text("prefix"),
690 )));
691 root.add_child(ParseTree::Rule(RuleNode::new(direct)));
692 root.add_child(ParseTree::Rule(RuleNode::new(wrapper)));
693
694 assert_eq!(root.child_count(), 3);
695 assert_eq!(root.child_rules(1).count(), 1);
696 assert_eq!(
697 root.child_rule(1).map(ParserRuleContext::text),
698 Some("direct".to_owned())
699 );
700 assert_eq!(
701 root.child_rule(2).map(ParserRuleContext::rule_index),
702 Some(2)
703 );
704 assert!(root.child_rule(3).is_none());
705 }
706
707 #[test]
708 fn finds_direct_terminal_children_by_token_type() {
709 let mut nested = ParserRuleContext::new(1, -1);
710 nested.add_child(ParseTree::Terminal(TerminalNode::new(
711 CommonToken::new(13).with_text("nested"),
712 )));
713
714 let mut root = ParserRuleContext::new(0, -1);
715 root.add_child(ParseTree::Error(ErrorNode::new(
716 CommonToken::new(12).with_text("error"),
717 )));
718 root.add_child(ParseTree::Terminal(TerminalNode::new(
719 CommonToken::new(10).with_text("direct"),
720 )));
721 root.add_child(ParseTree::Terminal(TerminalNode::new(
722 CommonToken::new(11).with_text("other"),
723 )));
724 root.add_child(ParseTree::Rule(RuleNode::new(nested)));
725
726 assert_eq!(root.child_count(), 4);
727 assert!(root.has_token(10));
728 assert!(root.has_token(12));
729 assert_eq!(
730 root.child_token(10).map(TerminalNode::text),
731 Some("direct".to_owned())
732 );
733 assert_eq!(
734 root.child_token(11).map(TerminalNode::text),
735 Some("other".to_owned())
736 );
737 assert_eq!(
738 root.child_token(12).map(TerminalNode::text),
739 Some("error".to_owned())
740 );
741 assert!(root.child_token(13).is_none());
742 }
743}