antlr4_runtime/
tree.rs

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

antlr4_runtime/
tree.rs
