Repository navigation
Expand file tree
/
Copy pathnodes.rs
More file actions
136 lines (117 loc) · 3.92 KB
/
Copy pathnodes.rs
File metadata and controls
136 lines (117 loc) · 3.92 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
use std::ops::Index;
use ruff_index::{IndexVec, newtype_index};
use ruff_python_ast::{Expr, Stmt};
use ruff_text_size::{Ranged, TextRange};
use crate::BranchId;
/// Id uniquely identifying an AST node in a program.
///
/// Using a `u32` is sufficient because Ruff only supports parsing documents with a size of max
/// `u32::max` and it is impossible to have more nodes than characters in the file. We use a
/// `NonZeroU32` to take advantage of memory layout optimizations.
#[newtype_index]
#[derive(Ord, PartialOrd)]
pub struct NodeId;
/// An AST node in a program, along with a pointer to its parent node (if any).
#[derive(Debug)]
struct NodeWithParent<'a> {
/// A pointer to the AST node.
node: NodeRef<'a>,
/// The ID of the parent of this node, if any.
parent: Option<NodeId>,
/// The branch ID of this node, if any.
branch: Option<BranchId>,
}
/// The nodes of a program indexed by [`NodeId`]
#[derive(Debug, Default)]
pub struct Nodes<'a> {
nodes: IndexVec<NodeId, NodeWithParent<'a>>,
}
impl<'a> Nodes<'a> {
/// Inserts a new AST node into the tree and returns its unique ID.
pub(crate) fn insert(
&mut self,
node: NodeRef<'a>,
parent: Option<NodeId>,
branch: Option<BranchId>,
) -> NodeId {
self.nodes.push(NodeWithParent {
node,
parent,
branch,
})
}
/// Return the [`NodeId`] of the parent node.
#[inline]
pub(crate) fn parent_id(&self, node_id: NodeId) -> Option<NodeId> {
self.nodes[node_id].parent
}
/// Return the [`BranchId`] of the branch node.
#[inline]
pub(crate) fn branch_id(&self, node_id: NodeId) -> Option<BranchId> {
self.nodes[node_id].branch
}
/// Returns an iterator over all [`NodeId`] ancestors, starting from the given [`NodeId`].
pub(crate) fn ancestor_ids(&self, node_id: NodeId) -> impl Iterator<Item = NodeId> + '_ {
std::iter::successors(Some(node_id), |&node_id| self.nodes[node_id].parent)
}
}
impl<'a> Index<NodeId> for Nodes<'a> {
type Output = NodeRef<'a>;
#[inline]
fn index(&self, index: NodeId) -> &Self::Output {
&self.nodes[index].node
}
}
/// A reference to an AST node. Like [`ruff_python_ast::AnyNodeRef`], but wraps the node
/// itself (like [`Stmt`]) rather than the narrowed type (like [`ruff_python_ast::StmtAssign`]).
///
/// TODO(charlie): Replace with [`ruff_python_ast::AnyNodeRef`]. This requires migrating
/// the rest of the codebase to use [`ruff_python_ast::AnyNodeRef`] and related abstractions,
/// like [`ruff_python_ast::ExprRef`] instead of [`Expr`].
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum NodeRef<'a> {
Stmt(&'a Stmt),
Expr(&'a Expr),
}
impl<'a> NodeRef<'a> {
/// Returns the [`Stmt`] if this is a statement, or `None` if the reference is to another
/// kind of AST node.
pub(crate) fn as_statement(&self) -> Option<&'a Stmt> {
match self {
NodeRef::Stmt(stmt) => Some(stmt),
NodeRef::Expr(_) => None,
}
}
/// Returns the [`Expr`] if this is a expression, or `None` if the reference is to another
/// kind of AST node.
pub(crate) fn as_expression(&self) -> Option<&'a Expr> {
match self {
NodeRef::Stmt(_) => None,
NodeRef::Expr(expr) => Some(expr),
}
}
pub(crate) fn is_statement(&self) -> bool {
self.as_statement().is_some()
}
pub(crate) fn is_expression(&self) -> bool {
self.as_expression().is_some()
}
}
impl Ranged for NodeRef<'_> {
fn range(&self) -> TextRange {
match self {
NodeRef::Stmt(stmt) => stmt.range(),
NodeRef::Expr(expr) => expr.range(),
}
}
}
impl<'a> From<&'a Expr> for NodeRef<'a> {
fn from(expr: &'a Expr) -> Self {
NodeRef::Expr(expr)
}
}
impl<'a> From<&'a Stmt> for NodeRef<'a> {
fn from(stmt: &'a Stmt) -> Self {
NodeRef::Stmt(stmt)
}
}