ruby_prism/
node_ext.rs

1//! Node extension methods for the prism parser.
2//!
3//! This module provides convenience methods on AST nodes that aren't generated
4//! from the config, mirroring Ruby's `node_ext.rb`.
5
6use std::fmt;
7
8use crate::{ConstantPathNode, ConstantPathTargetNode, ConstantReadNode, ConstantTargetNode, ConstantWriteNode, Node};
9
10/// Errors for constant path name computation.
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub enum ConstantPathError {
13    /// The constant path contains dynamic parts (e.g., `var::Bar::Baz`).
14    DynamicParts,
15    /// The constant path contains missing nodes (e.g., `Foo::`).
16    MissingNodes,
17}
18
19impl fmt::Display for ConstantPathError {
20    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
21        match self {
22            Self::DynamicParts => {
23                write!(f, "Constant path contains dynamic parts. Cannot compute full name")
24            },
25            Self::MissingNodes => {
26                write!(f, "Constant path contains missing nodes. Cannot compute full name")
27            },
28        }
29    }
30}
31
32impl std::error::Error for ConstantPathError {}
33
34/// Trait for nodes that can compute their full constant name.
35///
36/// Implemented by constant-related nodes (`ConstantReadNode`,
37/// `ConstantWriteNode`, `ConstantTargetNode`, `ConstantPathNode`, and
38/// `ConstantPathTargetNode`).
39pub trait FullName<'pr> {
40    /// Returns the list of parts for the full name of this constant.
41    ///
42    /// # Errors
43    ///
44    /// Returns [`ConstantPathError`] if the path contains dynamic parts or
45    /// missing nodes.
46    fn full_name_parts(&self) -> Result<Vec<&'pr [u8]>, ConstantPathError>;
47
48    /// Returns the full name of this constant.
49    ///
50    /// # Errors
51    ///
52    /// Returns [`ConstantPathError`] if the path contains dynamic parts or
53    /// missing nodes.
54    fn full_name(&self) -> Result<Vec<u8>, ConstantPathError> {
55        let parts = self.full_name_parts()?;
56        let mut result = Vec::new();
57        for (index, part) in parts.iter().enumerate() {
58            if index > 0 {
59                result.extend_from_slice(b"::");
60            }
61            result.extend_from_slice(part);
62        }
63        Ok(result)
64    }
65}
66
67/// Computes `full_name_parts` for a `Node` by dispatching to the appropriate
68/// `FullName` implementation.
69#[allow(clippy::option_if_let_else)]
70fn full_name_parts_for_node<'pr>(node: &Node<'pr>) -> Result<Vec<&'pr [u8]>, ConstantPathError> {
71    if let Some(path_node) = node.as_constant_path_node() {
72        path_node.full_name_parts()
73    } else if let Some(read_node) = node.as_constant_read_node() {
74        read_node.full_name_parts()
75    } else {
76        Err(ConstantPathError::DynamicParts)
77    }
78}
79
80/// Computes `full_name_parts` for a constant path node given its name and
81/// parent.
82fn constant_path_full_name_parts<'pr>(name: Option<crate::ConstantId<'pr>>, parent: Option<Node<'pr>>) -> Result<Vec<&'pr [u8]>, ConstantPathError> {
83    let name = name.ok_or(ConstantPathError::MissingNodes)?;
84
85    let mut parts = match parent {
86        Some(parent) => full_name_parts_for_node(&parent)?,
87        None => vec![b"".as_slice()],
88    };
89
90    parts.push(name.as_slice());
91    Ok(parts)
92}
93
94/// Implements `FullName` for simple constant nodes that have a `name()` method
95/// returning a single constant identifier.
96macro_rules! impl_simple_full_name {
97    ($node_type:ident) => {
98        impl<'pr> FullName<'pr> for $node_type<'pr> {
99            fn full_name_parts(&self) -> Result<Vec<&'pr [u8]>, ConstantPathError> {
100                Ok(vec![self.name().as_slice()])
101            }
102        }
103    };
104}
105
106impl_simple_full_name!(ConstantReadNode);
107impl_simple_full_name!(ConstantWriteNode);
108impl_simple_full_name!(ConstantTargetNode);
109
110impl<'pr> FullName<'pr> for ConstantPathNode<'pr> {
111    fn full_name_parts(&self) -> Result<Vec<&'pr [u8]>, ConstantPathError> {
112        constant_path_full_name_parts(self.name(), self.parent())
113    }
114}
115
116impl<'pr> FullName<'pr> for ConstantPathTargetNode<'pr> {
117    fn full_name_parts(&self) -> Result<Vec<&'pr [u8]>, ConstantPathError> {
118        constant_path_full_name_parts(self.name(), self.parent())
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::{ConstantPathError, FullName};
125    use crate::parse;
126
127    #[test]
128    fn test_full_name_for_constant_read_node() {
129        let result = parse(b"Foo");
130        let node = result.node().as_program_node().unwrap().statements().body().iter().next().unwrap();
131        let constant = node.as_constant_read_node().unwrap();
132
133        assert_eq!(constant.full_name_parts().unwrap(), vec![b"Foo".as_slice()]);
134        assert_eq!(constant.full_name().unwrap(), b"Foo");
135    }
136
137    #[test]
138    fn test_full_name_for_constant_write_node() {
139        let result = parse(b"Foo = 1");
140        let node = result.node().as_program_node().unwrap().statements().body().iter().next().unwrap();
141        let constant = node.as_constant_write_node().unwrap();
142
143        assert_eq!(constant.full_name_parts().unwrap(), vec![b"Foo".as_slice()]);
144        assert_eq!(constant.full_name().unwrap(), b"Foo");
145    }
146
147    #[test]
148    fn test_full_name_for_constant_target_node() {
149        let result = parse(b"Foo, Bar = [1, 2]");
150        let node = result.node().as_program_node().unwrap().statements().body().iter().next().unwrap();
151        let multi_write = node.as_multi_write_node().unwrap();
152        let target = multi_write.lefts().iter().next().unwrap();
153        let constant = target.as_constant_target_node().unwrap();
154
155        assert_eq!(constant.full_name_parts().unwrap(), vec![b"Foo".as_slice()]);
156        assert_eq!(constant.full_name().unwrap(), b"Foo");
157    }
158
159    #[test]
160    fn test_full_name_for_constant_path() {
161        let result = parse(b"Foo::Bar");
162        let node = result.node().as_program_node().unwrap().statements().body().iter().next().unwrap();
163        let constant_path = node.as_constant_path_node().unwrap();
164
165        assert_eq!(constant_path.full_name_parts().unwrap(), vec![b"Foo".as_slice(), b"Bar".as_slice()]);
166        assert_eq!(constant_path.full_name().unwrap(), b"Foo::Bar");
167    }
168
169    #[test]
170    fn test_full_name_for_constant_path_with_stovetop() {
171        let result = parse(b"::Foo::Bar");
172        let node = result.node().as_program_node().unwrap().statements().body().iter().next().unwrap();
173        let constant_path = node.as_constant_path_node().unwrap();
174
175        assert_eq!(constant_path.full_name_parts().unwrap(), vec![b"".as_slice(), b"Foo".as_slice(), b"Bar".as_slice()]);
176        assert_eq!(constant_path.full_name().unwrap(), b"::Foo::Bar");
177    }
178
179    #[test]
180    fn test_full_name_for_constant_path_with_self() {
181        let source = r"
182self::
183  Bar::Baz::
184    Qux
185";
186        let result = parse(source.as_bytes());
187        let node = result.node().as_program_node().unwrap().statements().body().iter().next().unwrap();
188        let constant_path = node.as_constant_path_node().unwrap();
189
190        assert_eq!(constant_path.full_name().unwrap_err(), ConstantPathError::DynamicParts);
191    }
192
193    #[test]
194    fn test_full_name_for_constant_path_with_variable() {
195        let source = r"
196foo::
197  Bar::Baz::
198    Qux
199";
200        let result = parse(source.as_bytes());
201        let node = result.node().as_program_node().unwrap().statements().body().iter().next().unwrap();
202        let constant_path = node.as_constant_path_node().unwrap();
203
204        assert_eq!(constant_path.full_name().unwrap_err(), ConstantPathError::DynamicParts);
205    }
206
207    #[test]
208    fn test_full_name_for_constant_path_with_missing_name() {
209        let result = parse(b"Foo::");
210        let node = result.node().as_program_node().unwrap().statements().body().iter().next().unwrap();
211        let constant_path = node.as_constant_path_node().unwrap();
212
213        assert_eq!(constant_path.full_name().unwrap_err(), ConstantPathError::MissingNodes);
214    }
215
216    #[test]
217    fn test_full_name_for_constant_path_target() {
218        let result = parse(b"Foo::Bar, Baz = [1, 2]");
219        let node = result.node().as_program_node().unwrap().statements().body().iter().next().unwrap();
220        let multi_write = node.as_multi_write_node().unwrap();
221        let target = multi_write.lefts().iter().next().unwrap();
222        let constant_path = target.as_constant_path_target_node().unwrap();
223
224        assert_eq!(constant_path.full_name_parts().unwrap(), vec![b"Foo".as_slice(), b"Bar".as_slice()]);
225        assert_eq!(constant_path.full_name().unwrap(), b"Foo::Bar");
226    }
227
228    #[test]
229    fn test_full_name_for_constant_path_target_with_stovetop() {
230        let result = parse(b"::Foo, Bar = [1, 2]");
231        let node = result.node().as_program_node().unwrap().statements().body().iter().next().unwrap();
232        let multi_write = node.as_multi_write_node().unwrap();
233        let target = multi_write.lefts().iter().next().unwrap();
234        let constant_path = target.as_constant_path_target_node().unwrap();
235
236        assert_eq!(constant_path.full_name_parts().unwrap(), vec![b"".as_slice(), b"Foo".as_slice()]);
237        assert_eq!(constant_path.full_name().unwrap(), b"::Foo");
238    }
239
240    #[test]
241    fn test_full_name_for_constant_path_target_with_self() {
242        let result = parse(b"self::Foo, Bar = [1, 2]");
243        let node = result.node().as_program_node().unwrap().statements().body().iter().next().unwrap();
244        let multi_write = node.as_multi_write_node().unwrap();
245        let target = multi_write.lefts().iter().next().unwrap();
246        let constant_path = target.as_constant_path_target_node().unwrap();
247
248        assert_eq!(constant_path.full_name().unwrap_err(), ConstantPathError::DynamicParts);
249    }
250}