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