rustpython_ruff_python_ast/
find_node.rs

1use crate::AnyNodeRef;
2use crate::visitor::source_order::{SourceOrderVisitor, TraversalSignal, walk_node};
3use ruff_text_size::{Ranged, TextRange};
4use std::fmt;
5use std::fmt::Formatter;
6
7/// Returns the node with a minimal range that fully contains `range`.
8///
9/// If `range` is empty and falls within a parser *synthesized* node generated during error recovery,
10/// then the first node with the given range is returned.
11///
12/// ## Panics
13/// Panics if `range` is not contained within `root`.
14pub fn covering_node(root: AnyNodeRef, range: TextRange) -> CoveringNode {
15    struct Visitor<'a> {
16        range: TextRange,
17        found: bool,
18        ancestors: Vec<AnyNodeRef<'a>>,
19    }
20
21    impl<'a> SourceOrderVisitor<'a> for Visitor<'a> {
22        fn enter_node(&mut self, node: AnyNodeRef<'a>) -> TraversalSignal {
23            // If the node fully contains the range, than it is a possible match but traverse into its children
24            // to see if there's a node with a narrower range.
25            if !self.found && node.range().contains_range(self.range) {
26                self.ancestors.push(node);
27                TraversalSignal::Traverse
28            } else {
29                TraversalSignal::Skip
30            }
31        }
32
33        fn leave_node(&mut self, node: AnyNodeRef<'a>) {
34            if !self.found && self.ancestors.last() == Some(&node) {
35                self.found = true;
36            }
37        }
38    }
39
40    assert!(
41        root.range().contains_range(range),
42        "Range is not contained within root"
43    );
44
45    let mut visitor = Visitor {
46        range,
47        found: false,
48        ancestors: Vec::new(),
49    };
50
51    walk_node(&mut visitor, root);
52    CoveringNode::from_ancestors(visitor.ancestors)
53}
54
55/// The node with a minimal range that fully contains the search range.
56pub struct CoveringNode<'a> {
57    /// The covering node, along with all of its ancestors up to the
58    /// root. The root is always the first element and the covering
59    /// node found is always the last node. This sequence is guaranteed
60    /// to be non-empty.
61    nodes: Vec<AnyNodeRef<'a>>,
62}
63
64impl<'a> CoveringNode<'a> {
65    /// Creates a new `CoveringNode` from a list of ancestor nodes.
66    /// The ancestors should be ordered from root to the covering node.
67    pub fn from_ancestors(ancestors: Vec<AnyNodeRef<'a>>) -> Self {
68        Self { nodes: ancestors }
69    }
70
71    /// Returns the covering node found.
72    pub fn node(&self) -> AnyNodeRef<'a> {
73        *self
74            .nodes
75            .last()
76            .expect("`CoveringNode::nodes` should always be non-empty")
77    }
78
79    /// Returns the node's parent.
80    pub fn parent(&self) -> Option<AnyNodeRef<'a>> {
81        let penultimate = self.nodes.len().checked_sub(2)?;
82        self.nodes.get(penultimate).copied()
83    }
84
85    /// Finds the first node that fully covers the range and fulfills
86    /// the given predicate.
87    ///
88    /// The "first" here means that the node closest to a leaf is
89    /// returned.
90    pub fn find_first(mut self, f: impl Fn(AnyNodeRef<'a>) -> bool) -> Result<Self, Self> {
91        let Some(index) = self.find_first_index(f) else {
92            return Err(self);
93        };
94        self.nodes.truncate(index + 1);
95        Ok(self)
96    }
97
98    /// Finds the last node that fully covers the range and fulfills
99    /// the given predicate.
100    ///
101    /// The "last" here means that after finding the "first" such node,
102    /// the highest ancestor found satisfying the given predicate is
103    /// returned. Note that this is *not* the same as finding the node
104    /// closest to the root that satisfies the given predictate.
105    pub fn find_last(mut self, f: impl Fn(AnyNodeRef<'a>) -> bool) -> Result<Self, Self> {
106        let Some(mut index) = self.find_first_index(&f) else {
107            return Err(self);
108        };
109        while index > 0 && f(self.nodes[index - 1]) {
110            index -= 1;
111        }
112        self.nodes.truncate(index + 1);
113        Ok(self)
114    }
115
116    /// Returns an iterator over the ancestor nodes, starting with the node itself
117    /// and walking towards the root.
118    pub fn ancestors(&self) -> impl DoubleEndedIterator<Item = AnyNodeRef<'a>> + '_ {
119        self.nodes.iter().copied().rev()
120    }
121
122    /// Finds the index of the node that fully covers the range and
123    /// fulfills the given predicate.
124    ///
125    /// If there are no nodes matching the given predictate, then
126    /// `None` is returned.
127    fn find_first_index(&self, f: impl Fn(AnyNodeRef<'a>) -> bool) -> Option<usize> {
128        self.nodes.iter().rposition(|node| f(*node))
129    }
130}
131
132impl fmt::Debug for CoveringNode<'_> {
133    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
134        f.debug_tuple("CoveringNode").field(&self.node()).finish()
135    }
136}
rustpython_ruff_python_ast/find_node.rs

rustpython_ruff_python_ast/
find_node.rs
