Skip to main content

fluidattacks_blends_domain/query/
paths.rs

1use alloc::collections::{BTreeMap, BTreeSet};
2use alloc::vec::Vec;
3
4use crate::query::adj_ast;
5use crate::syntax::{SyntaxGraph, SyntaxKind, SyntaxNode};
6use crate::{CodeGraph, NodeId};
7
8#[must_use]
9pub const fn nodes_by_type(graph: &SyntaxGraph) -> &BTreeMap<SyntaxKind, Vec<NodeId>> {
10    &graph.nodes_by_type
11}
12
13#[must_use]
14pub fn get_nodes_by_path<G: CodeGraph>(graph: &G, n_id: NodeId, path: &[&str]) -> BTreeSet<NodeId> {
15    let Some((first, rest)) = path.split_first() else {
16        return BTreeSet::new();
17    };
18    let matches = adj_ast(graph, n_id, Some(1), &[*first]);
19    if rest.is_empty() {
20        matches.into_iter().collect()
21    } else {
22        matches
23            .into_iter()
24            .flat_map(|child| get_nodes_by_path(graph, child, rest))
25            .collect()
26    }
27}
28
29#[must_use]
30pub fn get_node_by_path<G: CodeGraph>(graph: &G, n_id: NodeId, path: &[&str]) -> Option<NodeId> {
31    get_nodes_by_path(graph, n_id, path).into_iter().next()
32}
33
34#[must_use]
35pub fn get_args_n_ids(graph: &SyntaxGraph, n_id: NodeId) -> Vec<NodeId> {
36    let Some(args_parent_n_id) = graph.nodes.get(&n_id).and_then(SyntaxNode::arguments_id) else {
37        return Vec::new();
38    };
39    adj_ast(graph, args_parent_n_id, Some(1), &[])
40        .into_iter()
41        .filter(|arg_id| {
42            graph
43                .nodes
44                .get(arg_id)
45                .is_some_and(|node| !matches!(node, SyntaxNode::Comment { .. }))
46        })
47        .collect()
48}
49
50#[must_use]
51pub fn get_n_arg(graph: &SyntaxGraph, n_id: NodeId, arg_idx: usize) -> Option<NodeId> {
52    get_args_n_ids(graph, n_id).into_iter().nth(arg_idx)
53}
54
55#[cfg(test)]
56mod tests {
57    use super::{get_args_n_ids, get_n_arg, get_node_by_path, get_nodes_by_path, nodes_by_type};
58    use crate::ast::{AstGraph, AstNode};
59    use crate::syntax::{SyntaxGraph, SyntaxKind, SyntaxNode};
60    use crate::NodeId;
61    use alloc::borrow::ToOwned;
62    use alloc::collections::BTreeSet;
63    use alloc::vec;
64    use alloc::vec::Vec;
65
66    fn sample() -> AstGraph {
67        let mut ast = AstGraph::new();
68        ast.add_node(
69            NodeId(1),
70            AstNode::new(1, 1, "class_declaration".to_owned()),
71        );
72        ast.add_node(NodeId(2), AstNode::new(1, 1, "base_list".to_owned()));
73        ast.add_node(NodeId(3), AstNode::new(1, 1, "identifier".to_owned()));
74        ast.add_node(NodeId(4), AstNode::new(1, 1, "modifier".to_owned()));
75        ast.add_node(NodeId(5), AstNode::new(1, 2, "identifier".to_owned()));
76        ast.add_edge(NodeId(1), NodeId(2), 0);
77        ast.add_edge(NodeId(1), NodeId(4), 1);
78        ast.add_edge(NodeId(2), NodeId(3), 0);
79        ast.add_edge(NodeId(2), NodeId(5), 1);
80        ast
81    }
82
83    #[test]
84    fn nodes_by_type_exposes_the_index_maintained_by_add_node() {
85        let mut graph = SyntaxGraph::new();
86        graph.add_node(NodeId(9), SyntaxNode::Break);
87        graph.add_node(NodeId(1), SyntaxNode::File);
88        graph.add_node(NodeId(2), SyntaxNode::Break);
89
90        assert_eq!(
91            nodes_by_type(&graph).get(&SyntaxKind::Break),
92            Some(&vec![NodeId(9), NodeId(2)])
93        );
94        assert_eq!(
95            nodes_by_type(&graph).get(&SyntaxKind::File),
96            Some(&vec![NodeId(1)])
97        );
98        assert_eq!(nodes_by_type(&graph).get(&SyntaxKind::If), None);
99        assert!(nodes_by_type(&SyntaxGraph::new()).is_empty());
100    }
101
102    #[test]
103    fn nodes_by_path_collect_every_match_of_the_last_step() {
104        let ast = sample();
105        assert_eq!(
106            get_nodes_by_path(&ast, NodeId(1), &["base_list", "identifier"]),
107            BTreeSet::from([NodeId(3), NodeId(5)])
108        );
109        assert_eq!(
110            get_nodes_by_path(&ast, NodeId(1), &["base_list"]),
111            BTreeSet::from([NodeId(2)])
112        );
113    }
114
115    #[test]
116    fn nodes_by_path_are_empty_when_the_path_breaks_or_is_empty() {
117        let ast = sample();
118        assert_eq!(
119            get_nodes_by_path(&ast, NodeId(1), &["base_list", "block"]),
120            BTreeSet::new()
121        );
122        assert_eq!(get_nodes_by_path(&ast, NodeId(1), &[]), BTreeSet::new());
123    }
124
125    #[test]
126    fn follows_a_multi_step_label_path() {
127        let ast = sample();
128        assert_eq!(
129            get_node_by_path(&ast, NodeId(1), &["base_list", "identifier"]),
130            Some(NodeId(3))
131        );
132    }
133
134    #[test]
135    fn returns_the_direct_child_for_a_single_step_path() {
136        let ast = sample();
137        assert_eq!(
138            get_node_by_path(&ast, NodeId(1), &["base_list"]),
139            Some(NodeId(2))
140        );
141    }
142
143    #[test]
144    fn is_none_when_the_path_breaks_or_is_empty() {
145        let ast = sample();
146        assert_eq!(
147            get_node_by_path(&ast, NodeId(1), &["base_list", "block"]),
148            None
149        );
150        assert_eq!(get_node_by_path(&ast, NodeId(1), &["missing"]), None);
151        assert_eq!(get_node_by_path(&ast, NodeId(1), &[]), None);
152    }
153
154    fn invocation(arguments_id: Option<NodeId>) -> SyntaxNode {
155        SyntaxNode::MethodInvocation {
156            expression: "run".to_owned(),
157            object: None,
158            symbol_scope: None,
159            expression_id: None,
160            arguments_id,
161            object_id: None,
162            block_id: None,
163            receiver_type_fqn: None,
164        }
165    }
166
167    fn literal(value: &str) -> SyntaxNode {
168        SyntaxNode::Literal {
169            value: value.to_owned(),
170            value_type: "string".to_owned(),
171        }
172    }
173
174    fn call_with_commented_args() -> SyntaxGraph {
175        let mut graph = SyntaxGraph::new();
176        graph.add_node(NodeId(1), invocation(Some(NodeId(2))));
177        graph.add_node(NodeId(2), SyntaxNode::ArgumentList);
178        graph.add_node(NodeId(3), literal("first"));
179        graph.add_node(
180            NodeId(4),
181            SyntaxNode::Comment {
182                comment: "// second arg".to_owned(),
183            },
184        );
185        graph.add_node(NodeId(5), literal("second"));
186        graph.add_ast_edge(NodeId(1), NodeId(2));
187        graph.add_ast_edge(NodeId(2), NodeId(3));
188        graph.add_ast_edge(NodeId(2), NodeId(4));
189        graph.add_ast_edge(NodeId(2), NodeId(5));
190        graph
191    }
192
193    #[test]
194    fn args_skip_comments_interleaved_in_the_argument_list() {
195        let graph = call_with_commented_args();
196        assert_eq!(get_args_n_ids(&graph, NodeId(1)), [NodeId(3), NodeId(5)]);
197    }
198
199    #[test]
200    fn args_are_empty_without_an_argument_list() {
201        let mut graph = SyntaxGraph::new();
202        graph.add_node(NodeId(1), invocation(None));
203        graph.add_node(NodeId(2), invocation(Some(NodeId(9))));
204        graph.add_node(NodeId(3), literal("plain"));
205
206        assert_eq!(get_args_n_ids(&graph, NodeId(1)), Vec::new());
207        assert_eq!(get_args_n_ids(&graph, NodeId(2)), Vec::new());
208        assert_eq!(get_args_n_ids(&graph, NodeId(3)), Vec::new());
209        assert_eq!(get_args_n_ids(&graph, NodeId(404)), Vec::new());
210    }
211
212    #[test]
213    fn nth_arg_indexes_the_comment_free_positions() {
214        let graph = call_with_commented_args();
215        assert_eq!(get_n_arg(&graph, NodeId(1), 0), Some(NodeId(3)));
216        assert_eq!(get_n_arg(&graph, NodeId(1), 1), Some(NodeId(5)));
217        assert_eq!(get_n_arg(&graph, NodeId(1), 2), None);
218    }
219
220    #[test]
221    fn args_ignore_edges_pointing_at_unknown_nodes() {
222        let mut graph = call_with_commented_args();
223        graph.add_ast_edge(NodeId(2), NodeId(0));
224
225        assert_eq!(get_args_n_ids(&graph, NodeId(1)), [NodeId(3), NodeId(5)]);
226        assert_eq!(get_n_arg(&graph, NodeId(1), 0), Some(NodeId(3)));
227        assert_eq!(get_n_arg(&graph, NodeId(1), 1), Some(NodeId(5)));
228    }
229}