fluidattacks_blends_domain/query/
traversal.rs

1//! Counterpart of `blends/query/traversal.py`.
2
3use alloc::collections::BTreeSet;
4use alloc::vec::Vec;
5
6use crate::{CodeGraph, NodeId};
7
8/// Which edges a traversal follows.
9#[derive(Clone, Copy, PartialEq, Eq, Debug)]
10pub enum EdgeKind {
11    Any,
12    Ast,
13    Cfg,
14}
15
16fn children_for<G: CodeGraph>(graph: &G, n_id: NodeId, edges: EdgeKind) -> Vec<NodeId> {
17    match edges {
18        EdgeKind::Any => graph.children(n_id),
19        EdgeKind::Ast => graph.ast_children(n_id),
20        EdgeKind::Cfg => graph.cfg_children(n_id),
21    }
22}
23
24#[derive(Default)]
25struct AdjWalk {
26    nodes: Vec<NodeId>,
27    processed: BTreeSet<NodeId>,
28}
29
30fn collect_adj<G: CodeGraph>(
31    graph: &G,
32    n_id: NodeId,
33    depth: i64,
34    edges: EdgeKind,
35    walk: &mut AdjWalk,
36) {
37    if depth == 0 || !walk.processed.insert(n_id) {
38        return;
39    }
40    let childs = children_for(graph, n_id, edges);
41    walk.nodes.extend(&childs);
42    if depth == 1 {
43        return;
44    }
45    for c_id in childs {
46        collect_adj(graph, c_id, depth.saturating_sub(1), edges, walk);
47    }
48}
49
50pub fn adj_lazy<G: CodeGraph>(
51    graph: &G,
52    n_id: NodeId,
53    depth: Option<i64>,
54    edges: EdgeKind,
55) -> impl Iterator<Item = NodeId> {
56    let mut walk = AdjWalk::default();
57    collect_adj(graph, n_id, depth.unwrap_or(1), edges, &mut walk);
58    walk.nodes.into_iter()
59}
60
61#[must_use]
62pub fn adj<G: CodeGraph>(
63    graph: &G,
64    n_id: NodeId,
65    depth: Option<i64>,
66    edges: EdgeKind,
67) -> Vec<NodeId> {
68    adj_lazy(graph, n_id, depth, edges).collect()
69}
70
71#[must_use]
72pub fn adj_ast<G: CodeGraph>(
73    graph: &G,
74    n_id: NodeId,
75    depth: Option<i64>,
76    label_types: &[&str],
77) -> Vec<NodeId> {
78    filter_by_label_type(graph, adj(graph, n_id, depth, EdgeKind::Ast), label_types)
79}
80
81#[must_use]
82pub fn adj_cfg<G: CodeGraph>(
83    graph: &G,
84    n_id: NodeId,
85    depth: Option<i64>,
86    label_types: &[&str],
87) -> Vec<NodeId> {
88    filter_by_label_type(graph, adj(graph, n_id, depth, EdgeKind::Cfg), label_types)
89}
90
91fn parents_for<G: CodeGraph>(graph: &G, n_id: NodeId, edges: EdgeKind) -> Vec<NodeId> {
92    match edges {
93        EdgeKind::Any => graph.parents(n_id),
94        EdgeKind::Ast => graph.ast_parents(n_id),
95        EdgeKind::Cfg => graph.cfg_parents(n_id),
96    }
97}
98
99fn collect_pred<G: CodeGraph>(
100    graph: &G,
101    n_id: NodeId,
102    depth: i64,
103    edges: EdgeKind,
104    walk: &mut AdjWalk,
105) {
106    if depth == 0 || !walk.processed.insert(n_id) {
107        return;
108    }
109    let parents = parents_for(graph, n_id, edges);
110    walk.nodes.extend(&parents);
111    if depth == 1 {
112        return;
113    }
114    for p_id in parents {
115        collect_pred(graph, p_id, depth.saturating_sub(1), edges, walk);
116    }
117}
118
119pub fn pred_lazy<G: CodeGraph>(
120    graph: &G,
121    n_id: NodeId,
122    depth: Option<i64>,
123    edges: EdgeKind,
124) -> impl Iterator<Item = NodeId> {
125    let mut walk = AdjWalk::default();
126    collect_pred(graph, n_id, depth.unwrap_or(1), edges, &mut walk);
127    walk.nodes.into_iter()
128}
129
130#[must_use]
131pub fn pred<G: CodeGraph>(
132    graph: &G,
133    n_id: NodeId,
134    depth: Option<i64>,
135    edges: EdgeKind,
136) -> Vec<NodeId> {
137    pred_lazy(graph, n_id, depth, edges).collect()
138}
139
140#[must_use]
141pub fn pred_ast<G: CodeGraph>(graph: &G, n_id: NodeId, depth: Option<i64>) -> Vec<NodeId> {
142    pred(graph, n_id, depth, EdgeKind::Ast)
143}
144
145#[must_use]
146pub fn pred_cfg<G: CodeGraph>(graph: &G, n_id: NodeId, depth: Option<i64>) -> Vec<NodeId> {
147    pred(graph, n_id, depth, EdgeKind::Cfg)
148}
149
150fn filter_by_label_type<G: CodeGraph>(
151    graph: &G,
152    childs: Vec<NodeId>,
153    label_types: &[&str],
154) -> Vec<NodeId> {
155    if label_types.is_empty() {
156        return childs;
157    }
158    childs
159        .into_iter()
160        .filter(|c_id| {
161            graph
162                .label_type(*c_id)
163                .is_some_and(|kind| label_types.contains(&kind))
164        })
165        .collect()
166}
167
168#[cfg(test)]
169mod tests {
170    use super::{adj, adj_ast, adj_cfg, adj_lazy, pred, pred_ast, pred_cfg, pred_lazy, EdgeKind};
171    use crate::ast::{AstGraph, AstNode};
172    use crate::syntax::{SyntaxGraph, SyntaxNode};
173    use crate::NodeId;
174    use alloc::borrow::ToOwned;
175    use alloc::vec::Vec;
176
177    fn sample_ast() -> AstGraph {
178        let mut ast = AstGraph::new();
179        ast.add_node(NodeId(1), AstNode::new(1, 1, "stream".to_owned()));
180        ast.add_node(NodeId(2), AstNode::new(1, 1, "document".to_owned()));
181        ast.add_node(NodeId(3), AstNode::new(2, 1, "document".to_owned()));
182        ast.add_node(NodeId(4), AstNode::new(2, 1, "block_node".to_owned()));
183        ast.add_edge(NodeId(1), NodeId(2), 0);
184        ast.add_edge(NodeId(1), NodeId(3), 1);
185        ast.add_edge(NodeId(3), NodeId(4), 0);
186        ast
187    }
188
189    #[test]
190    fn adj_defaults_to_direct_children_in_id_order() {
191        let ast = sample_ast();
192        assert_eq!(
193            adj(&ast, NodeId(1), None, EdgeKind::Any),
194            [NodeId(2), NodeId(3)]
195        );
196        assert_eq!(adj(&ast, NodeId(4), None, EdgeKind::Any), []);
197        assert_eq!(adj(&ast, NodeId(1), Some(0), EdgeKind::Any), []);
198    }
199
200    #[test]
201    fn adj_depth_covers_descendants_and_minus_one_is_infinite() {
202        let ast = sample_ast();
203        assert_eq!(
204            adj(&ast, NodeId(1), Some(2), EdgeKind::Any),
205            [NodeId(2), NodeId(3), NodeId(4)]
206        );
207        assert_eq!(
208            adj(&ast, NodeId(1), Some(-1), EdgeKind::Any),
209            [NodeId(2), NodeId(3), NodeId(4)]
210        );
211    }
212
213    #[test]
214    fn pred_depths_mirror_the_python_comprehensive_suite() {
215        let mut ast = AstGraph::new();
216        for (n_id, kind) in [(1, "a"), (2, "b"), (3, "c"), (4, "d"), (5, "e"), (6, "f")] {
217            ast.add_node(NodeId(n_id), AstNode::new(1, 1, kind.to_owned()));
218        }
219        ast.add_edge(NodeId(1), NodeId(2), 0);
220        ast.add_edge(NodeId(1), NodeId(3), 1);
221        ast.add_edge(NodeId(2), NodeId(6), 0);
222        ast.add_edge(NodeId(3), NodeId(4), 0);
223        ast.add_edge(NodeId(4), NodeId(5), 0);
224
225        assert_eq!(
226            pred(&ast, NodeId(5), Some(-1), EdgeKind::Any),
227            [NodeId(4), NodeId(3), NodeId(1)]
228        );
229        assert_eq!(pred(&ast, NodeId(5), Some(0), EdgeKind::Any), []);
230        assert_eq!(pred(&ast, NodeId(5), None, EdgeKind::Any), [NodeId(4)]);
231        assert_eq!(
232            pred(&ast, NodeId(5), Some(2), EdgeKind::Any),
233            [NodeId(4), NodeId(3)]
234        );
235        assert_eq!(pred(&ast, NodeId(1), Some(-1), EdgeKind::Any), []);
236    }
237
238    #[test]
239    fn pred_lazy_yields_parents_in_order_expanding_each_node_once() {
240        let mut ast = AstGraph::new();
241        ast.add_node(NodeId(1), AstNode::new(1, 1, "a".to_owned()));
242        ast.add_node(NodeId(2), AstNode::new(1, 1, "b".to_owned()));
243        ast.add_node(NodeId(3), AstNode::new(1, 1, "c".to_owned()));
244        ast.add_edge(NodeId(1), NodeId(2), 0);
245        ast.add_edge(NodeId(2), NodeId(3), 0);
246        ast.add_edge(NodeId(1), NodeId(3), 1);
247
248        let direct: Vec<NodeId> = pred_lazy(&ast, NodeId(3), Some(1), EdgeKind::Any).collect();
249        assert_eq!(direct, [NodeId(1), NodeId(2)]);
250
251        let deep: Vec<NodeId> = pred_lazy(&ast, NodeId(3), Some(-1), EdgeKind::Any).collect();
252        assert_eq!(deep, [NodeId(1), NodeId(2), NodeId(1)]);
253
254        let zero: Vec<NodeId> = pred_lazy(&ast, NodeId(3), Some(0), EdgeKind::Any).collect();
255        assert_eq!(zero, []);
256    }
257
258    #[test]
259    fn pred_terminates_on_an_edge_cycle() {
260        let mut cyclic = AstGraph::new();
261        cyclic.add_node(NodeId(1), AstNode::new(1, 1, "a".to_owned()));
262        cyclic.add_node(NodeId(2), AstNode::new(1, 1, "b".to_owned()));
263        cyclic.add_edge(NodeId(1), NodeId(2), 0);
264        cyclic.add_edge(NodeId(2), NodeId(1), 0);
265
266        assert_eq!(
267            pred_ast(&cyclic, NodeId(1), Some(-1)),
268            [NodeId(2), NodeId(1)]
269        );
270        assert_eq!(pred_ast(&cyclic, NodeId(1), None), [NodeId(2)]);
271    }
272
273    #[test]
274    fn pred_follows_only_the_requested_edge_kind() {
275        let mut syntax = SyntaxGraph::new();
276        syntax.add_node(NodeId(1), SyntaxNode::File);
277        syntax.add_node(NodeId(2), SyntaxNode::ExecutionBlock);
278        syntax.add_node(NodeId(3), SyntaxNode::ExecutionBlock);
279        syntax.add_ast_edge(NodeId(1), NodeId(3));
280        syntax.add_cfg_edge(NodeId(2), NodeId(3));
281
282        assert_eq!(pred_ast(&syntax, NodeId(3), None), [NodeId(1)]);
283        assert_eq!(pred_cfg(&syntax, NodeId(3), None), [NodeId(2)]);
284        assert_eq!(
285            pred(&syntax, NodeId(3), None, EdgeKind::Any),
286            [NodeId(1), NodeId(2)]
287        );
288    }
289
290    #[test]
291    fn adj_lazy_repeats_a_shared_child_across_parents_expanding_it_once() {
292        let mut ast = AstGraph::new();
293        ast.add_node(NodeId(1), AstNode::new(1, 1, "a".to_owned()));
294        ast.add_node(NodeId(2), AstNode::new(1, 1, "b".to_owned()));
295        ast.add_node(NodeId(3), AstNode::new(1, 1, "c".to_owned()));
296        ast.add_edge(NodeId(1), NodeId(2), 0);
297        ast.add_edge(NodeId(1), NodeId(3), 1);
298        ast.add_edge(NodeId(2), NodeId(3), 0);
299
300        let yielded: Vec<NodeId> = adj_lazy(&ast, NodeId(1), Some(-1), EdgeKind::Any).collect();
301        assert_eq!(yielded, [NodeId(2), NodeId(3), NodeId(3)]);
302    }
303
304    #[test]
305    fn adj_terminates_on_an_edge_cycle() {
306        let mut ast = AstGraph::new();
307        ast.add_node(NodeId(1), AstNode::new(1, 1, "a".to_owned()));
308        ast.add_node(NodeId(2), AstNode::new(1, 1, "b".to_owned()));
309        ast.add_edge(NodeId(1), NodeId(2), 0);
310        ast.add_edge(NodeId(2), NodeId(1), 0);
311
312        assert_eq!(
313            adj(&ast, NodeId(1), Some(-1), EdgeKind::Any),
314            [NodeId(2), NodeId(1)]
315        );
316        assert_eq!(
317            adj(&ast, NodeId(2), Some(-1), EdgeKind::Any),
318            [NodeId(1), NodeId(2)]
319        );
320    }
321
322    #[test]
323    fn adj_ast_delegates_to_adj_and_filters_by_label_type() {
324        let ast = sample_ast();
325        assert_eq!(adj_ast(&ast, NodeId(1), None, &[]), [NodeId(2), NodeId(3)]);
326        assert_eq!(
327            adj_ast(&ast, NodeId(1), None, &["document"]),
328            [NodeId(2), NodeId(3)]
329        );
330        assert_eq!(adj_ast(&ast, NodeId(1), None, &["block_node"]), []);
331        assert_eq!(
332            adj_ast(&ast, NodeId(1), Some(-1), &["block_node"]),
333            [NodeId(4)]
334        );
335    }
336
337    fn sample_syntax() -> SyntaxGraph {
338        let mut graph = SyntaxGraph::new();
339        graph.add_node(NodeId(1), SyntaxNode::File);
340        graph.add_node(NodeId(2), SyntaxNode::ArrayInitializer);
341        graph.add_node(
342            NodeId(3),
343            SyntaxNode::Object {
344                name: None,
345                tf_reference: None,
346            },
347        );
348        graph.add_node(
349            NodeId(4),
350            SyntaxNode::Literal {
351                value: "doe".to_owned(),
352                value_type: "string".to_owned(),
353            },
354        );
355        graph.add_ast_edge(NodeId(1), NodeId(2));
356        graph.add_ast_edge(NodeId(1), NodeId(3));
357        graph.add_cfg_edge(NodeId(1), NodeId(3));
358        graph.add_cfg_edge(NodeId(3), NodeId(4));
359        graph
360    }
361
362    #[test]
363    fn adj_cfg_follows_only_edges_with_the_cfg_mark() {
364        let syntax = sample_syntax();
365        assert_eq!(adj_cfg(&syntax, NodeId(1), None, &[]), [NodeId(3)]);
366        assert_eq!(
367            adj_cfg(&syntax, NodeId(1), Some(-1), &[]),
368            [NodeId(3), NodeId(4)]
369        );
370        assert_eq!(adj_cfg(&syntax, NodeId(2), None, &[]), []);
371    }
372
373    #[test]
374    fn adj_cfg_filters_by_label_type() {
375        let syntax = sample_syntax();
376        assert_eq!(
377            adj_cfg(&syntax, NodeId(1), Some(-1), &["Literal"]),
378            [NodeId(4)]
379        );
380        assert_eq!(adj_cfg(&syntax, NodeId(1), None, &["Literal"]), []);
381    }
382}
fluidattacks_blends_domain/query/traversal.rs

fluidattacks_blends_domain/query/
traversal.rs
