Skip to main content

fluidattacks_blends_domain/syntax/
graph.rs

1//! Syntax graph: node table plus AST-mirroring adjacency.
2
3use alloc::collections::BTreeMap;
4use alloc::string::String;
5use alloc::vec::Vec;
6
7use crate::syntax::{SyntaxEdge, SyntaxKind, SyntaxNode};
8use crate::{Ast, Cfg, CodeGraph, NodeId};
9
10#[derive(Clone, PartialEq, Eq, Debug)]
11pub struct SanitizationEvent {
12    pub node_id: NodeId,
13    pub kind: String,
14    pub subpath: Option<NodeId>,
15}
16
17#[derive(Clone, Copy, PartialEq, Eq, Debug)]
18pub enum UsageRole {
19    Receiver,
20    Argument,
21    MemberWrite,
22}
23
24#[derive(Clone, PartialEq, Eq, Debug)]
25pub struct UsageEvent {
26    pub node_id: NodeId,
27    pub role: UsageRole,
28    pub arg_index: Option<i64>,
29    pub subpath: Option<NodeId>,
30}
31
32#[derive(Clone, PartialEq, Eq, Debug, Default)]
33pub struct SyntaxGraph {
34    pub nodes: BTreeMap<NodeId, SyntaxNode>,
35    pub edges: BTreeMap<NodeId, BTreeMap<NodeId, SyntaxEdge>>,
36    pub symbol_index: BTreeMap<String, BTreeMap<NodeId, Vec<NodeId>>>,
37    pub scope_parent: BTreeMap<NodeId, NodeId>,
38    pub sanitization_index: BTreeMap<String, BTreeMap<NodeId, Vec<SanitizationEvent>>>,
39    pub usage_index: BTreeMap<String, BTreeMap<NodeId, Vec<UsageEvent>>>,
40    pub subpath_index: BTreeMap<NodeId, Option<NodeId>>,
41    /// Callee short name -> invocations of it, used to reach a method's call sites.
42    pub call_index: BTreeMap<String, Vec<NodeId>>,
43    pub nodes_by_type: BTreeMap<SyntaxKind, Vec<NodeId>>,
44}
45
46impl SyntaxGraph {
47    #[must_use]
48    pub fn new() -> Self {
49        Self::default()
50    }
51
52    pub fn add_node(&mut self, id: NodeId, node: SyntaxNode) {
53        self.nodes_by_type
54            .entry(SyntaxKind::from(&node))
55            .or_default()
56            .push(id);
57        self.nodes.insert(id, node);
58    }
59
60    pub fn register_sanitization(&mut self, symbol: &str, scope: NodeId, event: SanitizationEvent) {
61        self.sanitization_index
62            .entry(String::from(symbol))
63            .or_default()
64            .entry(scope)
65            .or_default()
66            .push(event);
67    }
68
69    pub fn register_call(&mut self, name: &str, invocation_id: NodeId) {
70        self.call_index
71            .entry(String::from(name))
72            .or_default()
73            .push(invocation_id);
74    }
75
76    pub fn register_usage(&mut self, symbol: &str, scope: NodeId, event: UsageEvent) {
77        self.usage_index
78            .entry(String::from(symbol))
79            .or_default()
80            .entry(scope)
81            .or_default()
82            .push(event);
83    }
84
85    pub fn add_ast_edge(&mut self, from: NodeId, to: NodeId) {
86        self.edges
87            .entry(from)
88            .or_default()
89            .entry(to)
90            .or_default()
91            .ast = Some(Ast);
92    }
93
94    pub fn add_cfg_edge(&mut self, from: NodeId, to: NodeId) {
95        self.edges
96            .entry(from)
97            .or_default()
98            .entry(to)
99            .or_default()
100            .cfg = Some(Cfg);
101    }
102
103    pub fn finalize_symbol_index(&mut self) {
104        for scope_definitions in self.symbol_index.values_mut() {
105            for definition_list in scope_definitions.values_mut() {
106                definition_list.reverse();
107            }
108        }
109    }
110
111    pub fn remove_edge(&mut self, from: NodeId, to: NodeId) {
112        if let Some(adjacent) = self.edges.get_mut(&from) {
113            adjacent.remove(&to);
114            if adjacent.is_empty() {
115                self.edges.remove(&from);
116            }
117        }
118    }
119}
120
121impl CodeGraph for SyntaxGraph {
122    fn children(&self, n_id: NodeId) -> Vec<NodeId> {
123        self.edges
124            .get(&n_id)
125            .map(|adjacent| adjacent.keys().copied().collect())
126            .unwrap_or_default()
127    }
128
129    fn ast_children(&self, n_id: NodeId) -> Vec<NodeId> {
130        self.edges
131            .get(&n_id)
132            .map(|adjacent| {
133                adjacent
134                    .iter()
135                    .filter(|(_, edge)| edge.ast.is_some())
136                    .map(|(to, _)| *to)
137                    .collect()
138            })
139            .unwrap_or_default()
140    }
141
142    fn parents(&self, n_id: NodeId) -> Vec<NodeId> {
143        self.edges
144            .iter()
145            .filter(|(_, adjacent)| adjacent.contains_key(&n_id))
146            .map(|(from, _)| *from)
147            .collect()
148    }
149
150    fn ast_parents(&self, n_id: NodeId) -> Vec<NodeId> {
151        self.edges
152            .iter()
153            .filter(|(_, adjacent)| adjacent.get(&n_id).is_some_and(|edge| edge.ast.is_some()))
154            .map(|(from, _)| *from)
155            .collect()
156    }
157
158    fn cfg_parents(&self, n_id: NodeId) -> Vec<NodeId> {
159        self.edges
160            .iter()
161            .filter(|(_, adjacent)| adjacent.get(&n_id).is_some_and(|edge| edge.cfg.is_some()))
162            .map(|(from, _)| *from)
163            .collect()
164    }
165
166    fn cfg_children(&self, n_id: NodeId) -> Vec<NodeId> {
167        self.edges
168            .get(&n_id)
169            .map(|adjacent| {
170                adjacent
171                    .iter()
172                    .filter(|(_, edge)| edge.cfg.is_some())
173                    .map(|(to, _)| *to)
174                    .collect()
175            })
176            .unwrap_or_default()
177    }
178
179    fn label_type(&self, n_id: NodeId) -> Option<&str> {
180        self.nodes.get(&n_id).map(SyntaxNode::label_type)
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use super::SyntaxGraph;
187    use crate::syntax::{SyntaxEdge, SyntaxKind, SyntaxNode};
188    use crate::{Ast, Cfg, CodeGraph, NodeId};
189    use alloc::borrow::ToOwned;
190
191    #[test]
192    fn stores_nodes_and_edges_by_id() {
193        let mut graph = SyntaxGraph::new();
194        graph.add_node(
195            NodeId(1),
196            SyntaxNode::MissingNode {
197                node_type: "stream".to_owned(),
198            },
199        );
200        graph.add_ast_edge(NodeId(1), NodeId(2));
201
202        assert_eq!(
203            graph.nodes.get(&NodeId(1)),
204            Some(&SyntaxNode::MissingNode {
205                node_type: "stream".to_owned()
206            })
207        );
208        assert_eq!(
209            graph
210                .edges
211                .get(&NodeId(1))
212                .and_then(|adjacent| adjacent.get(&NodeId(2))),
213            Some(&SyntaxEdge {
214                ast: Some(Ast),
215                cfg: None
216            })
217        );
218    }
219
220    #[test]
221    fn add_node_indexes_ids_by_syntax_type_in_insertion_order() {
222        let mut graph = SyntaxGraph::new();
223        graph.add_node(NodeId(9), SyntaxNode::Break);
224        graph.add_node(NodeId(1), SyntaxNode::File);
225        graph.add_node(NodeId(2), SyntaxNode::Break);
226
227        assert_eq!(
228            graph.nodes_by_type.get(&SyntaxKind::Break),
229            Some(&alloc::vec![NodeId(9), NodeId(2)])
230        );
231        assert_eq!(
232            graph.nodes_by_type.get(&SyntaxKind::File),
233            Some(&alloc::vec![NodeId(1)])
234        );
235        assert_eq!(graph.nodes_by_type.get(&SyntaxKind::If), None);
236    }
237
238    #[test]
239    fn ast_edge_merges_into_existing_edge() {
240        let mut graph = SyntaxGraph::new();
241        graph.add_ast_edge(NodeId(1), NodeId(2));
242        graph.add_ast_edge(NodeId(1), NodeId(2));
243
244        assert_eq!(
245            graph
246                .edges
247                .get(&NodeId(1))
248                .map(alloc::collections::BTreeMap::len),
249            Some(1)
250        );
251    }
252
253    #[test]
254    fn ast_children_skip_edges_without_the_ast_mark() {
255        let mut graph = SyntaxGraph::new();
256        graph.add_ast_edge(NodeId(1), NodeId(2));
257        graph.edges.entry(NodeId(1)).or_default().insert(
258            NodeId(3),
259            SyntaxEdge {
260                ast: None,
261                cfg: None,
262            },
263        );
264
265        assert_eq!(graph.ast_children(NodeId(1)), [NodeId(2)]);
266    }
267
268    #[test]
269    fn cfg_edge_merges_onto_the_ast_edge() {
270        let mut graph = SyntaxGraph::new();
271        graph.add_ast_edge(NodeId(1), NodeId(2));
272        graph.add_cfg_edge(NodeId(1), NodeId(2));
273
274        assert_eq!(
275            graph
276                .edges
277                .get(&NodeId(1))
278                .and_then(|adjacent| adjacent.get(&NodeId(2))),
279            Some(&SyntaxEdge {
280                ast: Some(Ast),
281                cfg: Some(Cfg)
282            })
283        );
284    }
285
286    #[test]
287    fn remove_edge_drops_the_edge_and_prunes_the_empty_entry() {
288        let mut graph = SyntaxGraph::new();
289        graph.add_ast_edge(NodeId(1), NodeId(2));
290        graph.add_ast_edge(NodeId(1), NodeId(3));
291        graph.add_cfg_edge(NodeId(1), NodeId(2));
292
293        graph.remove_edge(NodeId(1), NodeId(2));
294
295        assert_eq!(graph.children(NodeId(1)), [NodeId(3)]);
296        assert_eq!(graph.parents(NodeId(2)), []);
297
298        graph.remove_edge(NodeId(1), NodeId(3));
299
300        assert!(!graph.edges.contains_key(&NodeId(1)));
301    }
302
303    #[test]
304    fn finalize_symbol_index_puts_the_most_recent_definition_first() {
305        let mut graph = SyntaxGraph::new();
306        graph
307            .symbol_index
308            .entry("x".to_owned())
309            .or_default()
310            .entry(NodeId(1))
311            .or_default()
312            .extend([NodeId(2), NodeId(5), NodeId(9)]);
313
314        graph.finalize_symbol_index();
315
316        assert_eq!(
317            graph
318                .symbol_index
319                .get("x")
320                .and_then(|scopes| scopes.get(&NodeId(1))),
321            Some(&alloc::vec![NodeId(9), NodeId(5), NodeId(2)])
322        );
323    }
324
325    #[test]
326    fn remove_edge_ignores_a_missing_edge() {
327        let mut graph = SyntaxGraph::new();
328        graph.add_ast_edge(NodeId(1), NodeId(2));
329
330        graph.remove_edge(NodeId(7), NodeId(8));
331        graph.remove_edge(NodeId(1), NodeId(9));
332
333        assert_eq!(graph.children(NodeId(1)), [NodeId(2)]);
334    }
335
336    #[test]
337    fn cfg_children_skip_edges_without_the_cfg_mark() {
338        let mut graph = SyntaxGraph::new();
339        graph.add_ast_edge(NodeId(1), NodeId(2));
340        graph.add_ast_edge(NodeId(1), NodeId(3));
341        graph.add_cfg_edge(NodeId(1), NodeId(3));
342
343        assert_eq!(graph.cfg_children(NodeId(1)), [NodeId(3)]);
344        assert_eq!(graph.cfg_children(NodeId(9)), []);
345    }
346
347    #[test]
348    fn register_call_groups_invocations_by_name() {
349        let mut graph = SyntaxGraph::new();
350        graph.register_call("Render", NodeId(5));
351        graph.register_call("Render", NodeId(9));
352        assert_eq!(
353            graph.call_index.get("Render").map(Vec::as_slice),
354            Some([NodeId(5), NodeId(9)].as_slice())
355        );
356    }
357}