antlr4_runtime/
tree.rs

1use crate::errors::AntlrError;
2use std::rc::Rc;
3
4use crate::token::{CommonToken, Token, TokenRef};
5use std::collections::BTreeMap;
6
7#[derive(Clone, Debug, Eq, PartialEq)]
8pub enum ParseTree {
9    Rule(RuleNode),
10    Terminal(TerminalNode),
11    Error(ErrorNode),
12}
13
14impl ParseTree {
15    pub fn text(&self) -> String {
16        match self {
17            Self::Rule(rule) => rule.text(),
18            Self::Terminal(node) => node.text(),
19            Self::Error(node) => node.text(),
20        }
21    }
22
23    pub fn to_string_tree<S: AsRef<str>>(&self, rule_names: &[S]) -> String {
24        match self {
25            Self::Rule(rule) => rule.to_string_tree(rule_names),
26            Self::Terminal(node) => escape_tree_text(&node.text()),
27            Self::Error(node) => escape_tree_text(&node.text()),
28        }
29    }
30
31    /// Finds the first rule node with `rule_index` in a depth-first walk.
32    pub fn first_rule(&self, rule_index: usize) -> Option<&Self> {
33        match self {
34            Self::Rule(rule) => {
35                if rule.context().rule_index() == rule_index {
36                    return Some(self);
37                }
38                rule.context()
39                    .children()
40                    .iter()
41                    .find_map(|child| child.first_rule(rule_index))
42            }
43            Self::Terminal(_) | Self::Error(_) => None,
44        }
45    }
46
47    /// Finds the stop token for the first rule node with `rule_index`.
48    pub fn first_rule_stop(&self, rule_index: usize) -> Option<&CommonToken> {
49        let Self::Rule(rule) = self else {
50            return None;
51        };
52        if rule.context().rule_index() == rule_index {
53            return rule.context().stop();
54        }
55        rule.context()
56            .children()
57            .iter()
58            .find_map(|child| child.first_rule_stop(rule_index))
59    }
60
61    /// Reads an integer return value from the first rule node with
62    /// `rule_index`, matching ANTLR's `$label.value` resolution for labeled
63    /// rule references in the runtime testsuite.
64    pub fn first_rule_int_return(&self, rule_index: usize, name: &str) -> Option<i64> {
65        let Self::Rule(rule) = self else {
66            return None;
67        };
68        if rule.context().rule_index() == rule_index {
69            return rule.context().int_return(name);
70        }
71        rule.context()
72            .children()
73            .iter()
74            .find_map(|child| child.first_rule_int_return(rule_index, name))
75    }
76
77    /// Finds the first recovery error token in a depth-first walk.
78    pub fn first_error_token(&self) -> Option<&CommonToken> {
79        match self {
80            Self::Rule(rule) => rule
81                .context()
82                .children()
83                .iter()
84                .find_map(Self::first_error_token),
85            Self::Terminal(_) => None,
86            Self::Error(node) => Some(node.symbol()),
87        }
88    }
89
90    /// Returns the first rule invocation stack for `rule_index`, ordered from
91    /// the selected rule outward to the root rule.
92    pub fn rule_invocation_stack<S: AsRef<str>>(
93        &self,
94        rule_index: usize,
95        rule_names: &[S],
96    ) -> Option<Vec<String>> {
97        let mut stack = Vec::new();
98        if self.find_rule_path(rule_index, rule_names, &mut stack) {
99            stack.reverse();
100            return Some(stack);
101        }
102        None
103    }
104
105    fn find_rule_path<S: AsRef<str>>(
106        &self,
107        rule_index: usize,
108        rule_names: &[S],
109        stack: &mut Vec<String>,
110    ) -> bool {
111        let Self::Rule(rule) = self else {
112            return false;
113        };
114        let current_index = rule.context().rule_index();
115        stack.push(
116            rule_names
117                .get(current_index)
118                .map_or("<unknown>", |name| name.as_ref())
119                .to_owned(),
120        );
121        if current_index == rule_index
122            || rule
123                .context()
124                .children()
125                .iter()
126                .any(|child| child.find_rule_path(rule_index, rule_names, stack))
127        {
128            return true;
129        }
130        stack.pop();
131        false
132    }
133}
134
135fn escape_tree_text(text: &str) -> String {
136    let mut escaped = String::with_capacity(text.len());
137    for ch in text.chars() {
138        match ch {
139            '\n' => escaped.push_str("\\n"),
140            '\r' => escaped.push_str("\\r"),
141            '\t' => escaped.push_str("\\t"),
142            _ => escaped.push(ch),
143        }
144    }
145    escaped
146}
147
148#[derive(Clone, Debug, Eq, PartialEq)]
149pub struct RuleNode {
150    context: ParserRuleContext,
151}
152
153impl RuleNode {
154    pub const fn new(context: ParserRuleContext) -> Self {
155        Self { context }
156    }
157
158    pub const fn context(&self) -> &ParserRuleContext {
159        &self.context
160    }
161
162    pub const fn context_mut(&mut self) -> &mut ParserRuleContext {
163        &mut self.context
164    }
165
166    pub fn text(&self) -> String {
167        self.context.text()
168    }
169
170    pub fn to_string_tree<S: AsRef<str>>(&self, rule_names: &[S]) -> String {
171        self.context.to_string_tree(rule_names)
172    }
173}
174
175#[derive(Clone, Debug, Eq, PartialEq)]
176pub struct ParserRuleContext {
177    rule_index: usize,
178    invoking_state: isize,
179    alt_number: usize,
180    start: Option<TokenRef>,
181    stop: Option<TokenRef>,
182    int_returns: Option<Box<IntReturns>>,
183    children: Vec<ParseTree>,
184    /// Whether any child has been offered to this context, independent of whether
185    /// the tree was actually built. `children` stays empty when
186    /// `Parser::set_build_parse_trees(false)`, so generated recovery uses this
187    /// flag (not `children.is_empty()`) to tell whether the rule has matched
188    /// anything yet.
189    matched_child: bool,
190    // Boxed: an `AntlrError` is large and only set on the rare error path, so
191    // keeping it behind a pointer keeps `ParserRuleContext` (and thus the
192    // `ParseTree::Rule` variant) compact.
193    exception: Option<Box<AntlrError>>,
194}
195
196#[derive(Clone, Debug, Default, Eq, PartialEq)]
197struct IntReturns(BTreeMap<String, i64>);
198
199impl ParserRuleContext {
200    pub const fn new(rule_index: usize, invoking_state: isize) -> Self {
201        Self {
202            rule_index,
203            invoking_state,
204            alt_number: 0,
205            start: None,
206            stop: None,
207            int_returns: None,
208            children: Vec::new(),
209            matched_child: false,
210            exception: None,
211        }
212    }
213
214    pub(crate) fn with_child_capacity(
215        rule_index: usize,
216        invoking_state: isize,
217        capacity: usize,
218    ) -> Self {
219        Self {
220            rule_index,
221            invoking_state,
222            alt_number: 0,
223            start: None,
224            stop: None,
225            int_returns: None,
226            children: Vec::with_capacity(capacity),
227            matched_child: false,
228            exception: None,
229        }
230    }
231
232    pub const fn rule_index(&self) -> usize {
233        self.rule_index
234    }
235
236    pub const fn invoking_state(&self) -> isize {
237        self.invoking_state
238    }
239
240    pub const fn alt_number(&self) -> usize {
241        self.alt_number
242    }
243
244    pub const fn set_alt_number(&mut self, alt_number: usize) {
245        self.alt_number = alt_number;
246    }
247
248    pub fn start(&self) -> Option<&CommonToken> {
249        self.start.as_deref()
250    }
251
252    pub(crate) fn start_ref(&self) -> Option<TokenRef> {
253        self.start.as_ref().map(Rc::clone)
254    }
255
256    pub fn stop(&self) -> Option<&CommonToken> {
257        self.stop.as_deref()
258    }
259
260    pub fn set_start(&mut self, token: CommonToken) {
261        self.start = Some(Rc::new(token));
262    }
263
264    pub(crate) fn set_start_ref(&mut self, token: TokenRef) {
265        self.start = Some(token);
266    }
267
268    pub fn set_stop(&mut self, token: CommonToken) {
269        self.stop = Some(Rc::new(token));
270    }
271
272    pub(crate) fn set_stop_ref(&mut self, token: TokenRef) {
273        self.stop = Some(token);
274    }
275
276    /// Stores a generated integer return value on this rule context.
277    pub fn set_int_return(&mut self, name: impl Into<String>, value: i64) {
278        self.int_returns
279            .get_or_insert_with(Box::default)
280            .0
281            .insert(name.into(), value);
282    }
283
284    /// Reads a generated integer return value from this rule context.
285    pub fn int_return(&self, name: &str) -> Option<i64> {
286        self.int_returns
287            .as_ref()
288            .and_then(|values| values.0.get(name).copied())
289    }
290
291    pub fn exception(&self) -> Option<&AntlrError> {
292        self.exception.as_deref()
293    }
294
295    pub fn set_exception(&mut self, error: AntlrError) {
296        self.exception = Some(Box::new(error));
297    }
298
299    pub fn children(&self) -> &[ParseTree] {
300        &self.children
301    }
302
303    /// Returns the number of direct children in this context.
304    pub const fn child_count(&self) -> usize {
305        self.children.len()
306    }
307
308    /// Finds the first direct child rule with `rule_index`.
309    pub fn child_rule(&self, rule_index: usize) -> Option<&Self> {
310        self.child_rules(rule_index).next()
311    }
312
313    /// Iterates over direct child rules with `rule_index`.
314    pub fn child_rules(&self, rule_index: usize) -> impl Iterator<Item = &Self> + '_ {
315        self.children.iter().filter_map(move |child| match child {
316            ParseTree::Rule(rule) if rule.context().rule_index() == rule_index => {
317                Some(rule.context())
318            }
319            _ => None,
320        })
321    }
322
323    /// Finds the first direct token child with `token_type`.
324    ///
325    /// Includes recovery error nodes, which ANTLR treats as terminal nodes for
326    /// token-getter helpers.
327    pub fn child_token(&self, token_type: i32) -> Option<&TerminalNode> {
328        self.children.iter().find_map(|child| match child {
329            ParseTree::Terminal(node) if node.symbol().token_type() == token_type => Some(node),
330            ParseTree::Error(node) if node.symbol().token_type() == token_type => {
331                Some(node.terminal())
332            }
333            _ => None,
334        })
335    }
336
337    /// Returns whether this context has a direct token child with `token_type`.
338    pub fn has_token(&self, token_type: i32) -> bool {
339        self.child_token(token_type).is_some()
340    }
341
342    pub fn add_child(&mut self, child: ParseTree) {
343        self.matched_child = true;
344        self.children.push(child);
345    }
346
347    /// Records that a child was matched without storing it (used when parse-tree
348    /// construction is disabled). Keeps `has_matched_child` accurate even though
349    /// `children` stays empty.
350    pub const fn note_matched_child(&mut self) {
351        self.matched_child = true;
352    }
353
354    /// Whether this context has matched at least one child (token, rule, or error
355    /// node) so far, regardless of whether parse-tree construction is enabled.
356    pub const fn has_matched_child(&self) -> bool {
357        self.matched_child
358    }
359
360    pub fn text(&self) -> String {
361        self.children.iter().map(ParseTree::text).collect()
362    }
363
364    pub fn to_string_tree<S: AsRef<str>>(&self, rule_names: &[S]) -> String {
365        let name = rule_names
366            .get(self.rule_index)
367            .map_or("<unknown>", |name| name.as_ref());
368        let display_name = if self.alt_number == 0 {
369            name.to_owned()
370        } else {
371            format!("{name}:{}", self.alt_number)
372        };
373        if self.children.is_empty() {
374            return display_name;
375        }
376        let children = self
377            .children
378            .iter()
379            .map(|child| child.to_string_tree(rule_names))
380            .collect::<Vec<_>>()
381            .join(" ");
382        format!("({display_name} {children})")
383    }
384}
385
386#[derive(Clone, Debug, Eq, PartialEq)]
387pub struct TerminalNode {
388    token: TokenRef,
389}
390
391impl TerminalNode {
392    pub fn new(token: CommonToken) -> Self {
393        Self {
394            token: Rc::new(token),
395        }
396    }
397
398    pub(crate) const fn from_ref(token: TokenRef) -> Self {
399        Self { token }
400    }
401
402    pub fn symbol(&self) -> &CommonToken {
403        self.token.as_ref()
404    }
405
406    pub fn text(&self) -> String {
407        self.token.text().unwrap_or("").to_owned()
408    }
409}
410
411#[derive(Clone, Debug, Eq, PartialEq)]
412pub struct ErrorNode {
413    terminal: TerminalNode,
414}
415
416impl ErrorNode {
417    pub fn new(token: CommonToken) -> Self {
418        Self {
419            terminal: TerminalNode::new(token),
420        }
421    }
422
423    pub(crate) const fn from_ref(token: TokenRef) -> Self {
424        Self {
425            terminal: TerminalNode::from_ref(token),
426        }
427    }
428
429    const fn terminal(&self) -> &TerminalNode {
430        &self.terminal
431    }
432
433    pub fn symbol(&self) -> &CommonToken {
434        self.terminal.symbol()
435    }
436
437    pub fn text(&self) -> String {
438        self.terminal.text()
439    }
440}
441
442pub trait ParseTreeListener {
443    fn enter_every_rule(&mut self, _ctx: &ParserRuleContext) -> Result<(), AntlrError> {
444        Ok(())
445    }
446
447    fn exit_every_rule(&mut self, _ctx: &ParserRuleContext) -> Result<(), AntlrError> {
448        Ok(())
449    }
450
451    fn visit_terminal(&mut self, _node: &TerminalNode) -> Result<(), AntlrError> {
452        Ok(())
453    }
454
455    fn visit_error_node(&mut self, _node: &ErrorNode) -> Result<(), AntlrError> {
456        Ok(())
457    }
458}
459
460#[derive(Debug, Default)]
461pub struct ParseTreeWalker;
462
463impl ParseTreeWalker {
464    /// Walks a parse tree depth-first, invoking listener callbacks in ANTLR's
465    /// enter/child/exit order for rule nodes.
466    pub fn walk<L: ParseTreeListener>(
467        listener: &mut L,
468        tree: &ParseTree,
469    ) -> Result<(), AntlrError> {
470        match tree {
471            ParseTree::Rule(rule) => {
472                listener.enter_every_rule(rule.context())?;
473                for child in rule.context().children() {
474                    Self::walk(listener, child)?;
475                }
476                listener.exit_every_rule(rule.context())
477            }
478            ParseTree::Terminal(node) => listener.visit_terminal(node),
479            ParseTree::Error(node) => listener.visit_error_node(node),
480        }
481    }
482}
483
484#[cfg(test)]
485mod tests {
486    use super::*;
487    use crate::token::CommonToken;
488
489    #[test]
490    fn renders_rule_tree() {
491        let mut ctx = ParserRuleContext::new(0, -1);
492        ctx.add_child(ParseTree::Terminal(TerminalNode::new(
493            CommonToken::new(1).with_text("x"),
494        )));
495        let tree = ParseTree::Rule(RuleNode::new(ctx));
496        assert_eq!(tree.to_string_tree(&["expr"]), "(expr x)");
497    }
498
499    #[test]
500    fn finds_first_rule_depth_first() {
501        let mut nested = ParserRuleContext::new(1, -1);
502        nested.add_child(ParseTree::Terminal(TerminalNode::new(
503            CommonToken::new(1).with_text("x"),
504        )));
505
506        let mut root = ParserRuleContext::new(0, -1);
507        root.add_child(ParseTree::Rule(RuleNode::new(nested)));
508        let tree = ParseTree::Rule(RuleNode::new(root));
509
510        let rule = tree.first_rule(1).expect("nested rule should be found");
511        assert_eq!(
512            rule.to_string_tree(&["root".to_owned(), "child".to_owned()]),
513            "(child x)"
514        );
515        assert!(tree.first_rule(2).is_none());
516    }
517
518    #[test]
519    fn reports_rule_invocation_stack_from_leaf_to_root() {
520        let mut nested = ParserRuleContext::new(1, -1);
521        nested.add_child(ParseTree::Terminal(TerminalNode::new(
522            CommonToken::new(1).with_text("x"),
523        )));
524
525        let mut root = ParserRuleContext::new(0, -1);
526        root.add_child(ParseTree::Rule(RuleNode::new(nested)));
527        let tree = ParseTree::Rule(RuleNode::new(root));
528
529        assert_eq!(
530            tree.rule_invocation_stack(1, &["s", "a"]),
531            Some(vec!["a".to_owned(), "s".to_owned()])
532        );
533    }
534
535    #[test]
536    fn finds_direct_child_rules_by_index() {
537        let mut direct = ParserRuleContext::new(1, -1);
538        direct.add_child(ParseTree::Terminal(TerminalNode::new(
539            CommonToken::new(10).with_text("direct"),
540        )));
541
542        let mut nested_match = ParserRuleContext::new(1, -1);
543        nested_match.add_child(ParseTree::Terminal(TerminalNode::new(
544            CommonToken::new(11).with_text("nested"),
545        )));
546        let mut wrapper = ParserRuleContext::new(2, -1);
547        wrapper.add_child(ParseTree::Rule(RuleNode::new(nested_match)));
548
549        let mut root = ParserRuleContext::new(0, -1);
550        root.add_child(ParseTree::Terminal(TerminalNode::new(
551            CommonToken::new(12).with_text("prefix"),
552        )));
553        root.add_child(ParseTree::Rule(RuleNode::new(direct)));
554        root.add_child(ParseTree::Rule(RuleNode::new(wrapper)));
555
556        assert_eq!(root.child_count(), 3);
557        assert_eq!(root.child_rules(1).count(), 1);
558        assert_eq!(
559            root.child_rule(1).map(ParserRuleContext::text),
560            Some("direct".to_owned())
561        );
562        assert_eq!(
563            root.child_rule(2).map(ParserRuleContext::rule_index),
564            Some(2)
565        );
566        assert!(root.child_rule(3).is_none());
567    }
568
569    #[test]
570    fn finds_direct_terminal_children_by_token_type() {
571        let mut nested = ParserRuleContext::new(1, -1);
572        nested.add_child(ParseTree::Terminal(TerminalNode::new(
573            CommonToken::new(13).with_text("nested"),
574        )));
575
576        let mut root = ParserRuleContext::new(0, -1);
577        root.add_child(ParseTree::Error(ErrorNode::new(
578            CommonToken::new(12).with_text("error"),
579        )));
580        root.add_child(ParseTree::Terminal(TerminalNode::new(
581            CommonToken::new(10).with_text("direct"),
582        )));
583        root.add_child(ParseTree::Terminal(TerminalNode::new(
584            CommonToken::new(11).with_text("other"),
585        )));
586        root.add_child(ParseTree::Rule(RuleNode::new(nested)));
587
588        assert_eq!(root.child_count(), 4);
589        assert!(root.has_token(10));
590        assert!(root.has_token(12));
591        assert_eq!(
592            root.child_token(10).map(TerminalNode::text),
593            Some("direct".to_owned())
594        );
595        assert_eq!(
596            root.child_token(11).map(TerminalNode::text),
597            Some("other".to_owned())
598        );
599        assert_eq!(
600            root.child_token(12).map(TerminalNode::text),
601            Some("error".to_owned())
602        );
603        assert!(root.child_token(13).is_none());
604    }
605}
antlr4_runtime/tree.rs

antlr4_runtime/
tree.rs
