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/// A value sanitized while flowing through `node_id`, tracked for taint analysis.
11#[derive(Clone, PartialEq, Eq, Debug)]
12#[non_exhaustive]
13pub struct SanitizationEvent {
14    /// Id of the node performing the sanitization.
15    pub node_id: NodeId,
16    /// The sanitization's kind, e.g. the sanitizer function's name.
17    pub kind: String,
18    /// Id of the object path sanitized, when only part of a value is affected.
19    pub subpath: Option<NodeId>,
20}
21
22/// How a node participates in a usage event tracked by [`SyntaxGraph::usage_index`].
23#[derive(Clone, Copy, PartialEq, Eq, Debug)]
24#[allow(
25    clippy::exhaustive_enums,
26    reason = "a fixed ontology: the usage roles this graph's analyses distinguish"
27)]
28pub enum UsageRole {
29    /// The symbol is the receiver of a method call.
30    Receiver,
31    /// The symbol is passed as a call argument.
32    Argument,
33    /// The symbol is a member being written to.
34    MemberWrite,
35}
36
37/// A recorded use of a symbol, tracked for data-flow analysis.
38#[derive(Clone, PartialEq, Eq, Debug)]
39#[non_exhaustive]
40pub struct UsageEvent {
41    /// Id of the node where the symbol is used.
42    pub node_id: NodeId,
43    /// How the symbol participates at that node.
44    pub role: UsageRole,
45    /// The argument's position, for [`UsageRole::Argument`].
46    pub arg_index: Option<i64>,
47    /// Id of the object path used, when only part of a value is affected.
48    pub subpath: Option<NodeId>,
49}
50
51/// The syntax graph: every parsed node plus the AST and CFG edges between them.
52#[derive(Clone, PartialEq, Eq, Debug, Default)]
53#[non_exhaustive]
54pub struct SyntaxGraph {
55    /// Every node, keyed by id.
56    pub nodes: BTreeMap<NodeId, SyntaxNode>,
57    /// Adjacency: source id -> target id -> which edge kinds connect them.
58    pub edges: BTreeMap<NodeId, BTreeMap<NodeId, SyntaxEdge>>,
59    /// Symbol name -> scope id -> definition ids in that scope, most recent first.
60    pub symbol_index: BTreeMap<String, BTreeMap<NodeId, Vec<NodeId>>>,
61    /// Scope id -> its enclosing scope's id.
62    pub scope_parent: BTreeMap<NodeId, NodeId>,
63    /// Symbol name -> scope id -> sanitization events recorded for it.
64    pub sanitization_index: BTreeMap<String, BTreeMap<NodeId, Vec<SanitizationEvent>>>,
65    /// Symbol name -> scope id -> usage events recorded for it.
66    pub usage_index: BTreeMap<String, BTreeMap<NodeId, Vec<UsageEvent>>>,
67    /// Node id -> the object path id it was built from, when it has one.
68    pub subpath_index: BTreeMap<NodeId, Option<NodeId>>,
69    /// Callee short name -> invocations of it, used to reach a method's call sites.
70    pub call_index: BTreeMap<String, Vec<NodeId>>,
71    /// Syntax kind -> ids of every node of that kind, in insertion order.
72    pub nodes_by_type: BTreeMap<SyntaxKind, Vec<NodeId>>,
73}
74
75impl SyntaxGraph {
76    /// An empty syntax graph.
77    #[must_use]
78    pub fn new() -> Self {
79        Self::default()
80    }
81
82    /// Inserts `node` under `id`, indexing it by its syntax kind.
83    pub fn add_node(&mut self, id: NodeId, node: SyntaxNode) {
84        self.nodes_by_type
85            .entry(SyntaxKind::from(&node))
86            .or_default()
87            .push(id);
88        self.nodes.insert(id, node);
89    }
90
91    /// Records a sanitization of `symbol` as seen within `scope`.
92    pub fn register_sanitization(&mut self, symbol: &str, scope: NodeId, event: SanitizationEvent) {
93        self.sanitization_index
94            .entry(String::from(symbol))
95            .or_default()
96            .entry(scope)
97            .or_default()
98            .push(event);
99    }
100
101    /// Records that `invocation_id` calls a callee named `name`.
102    pub fn register_call(&mut self, name: &str, invocation_id: NodeId) {
103        self.call_index
104            .entry(String::from(name))
105            .or_default()
106            .push(invocation_id);
107    }
108
109    /// Records a use of `symbol` as seen within `scope`.
110    pub fn register_usage(&mut self, symbol: &str, scope: NodeId, event: UsageEvent) {
111        self.usage_index
112            .entry(String::from(symbol))
113            .or_default()
114            .entry(scope)
115            .or_default()
116            .push(event);
117    }
118
119    /// Connects `from` to `to` with an AST edge, merging into any existing edge.
120    pub fn add_ast_edge(&mut self, from: NodeId, to: NodeId) {
121        self.edges
122            .entry(from)
123            .or_default()
124            .entry(to)
125            .or_default()
126            .ast = Some(Ast);
127    }
128
129    /// Connects `from` to `to` with a CFG edge, merging into any existing edge.
130    pub fn add_cfg_edge(&mut self, from: NodeId, to: NodeId) {
131        self.edges
132            .entry(from)
133            .or_default()
134            .entry(to)
135            .or_default()
136            .cfg = Some(Cfg);
137    }
138
139    /// Reverses each scope's definition list so lookups see the most recent first.
140    ///
141    /// Call once after the graph is fully built; definitions are appended in
142    /// parse order, which is the opposite of the shadowing order callers need.
143    pub fn finalize_symbol_index(&mut self) {
144        for scope_definitions in self.symbol_index.values_mut() {
145            for definition_list in scope_definitions.values_mut() {
146                definition_list.reverse();
147            }
148        }
149    }
150
151    /// Removes the edge from `from` to `to`, if one exists.
152    pub fn remove_edge(&mut self, from: NodeId, to: NodeId) {
153        if let Some(adjacent) = self.edges.get_mut(&from) {
154            adjacent.remove(&to);
155            if adjacent.is_empty() {
156                self.edges.remove(&from);
157            }
158        }
159    }
160}
161
162impl CodeGraph for SyntaxGraph {
163    fn children(&self, n_id: NodeId) -> Vec<NodeId> {
164        self.edges
165            .get(&n_id)
166            .map(|adjacent| adjacent.keys().copied().collect())
167            .unwrap_or_default()
168    }
169
170    fn ast_children(&self, n_id: NodeId) -> Vec<NodeId> {
171        self.edges
172            .get(&n_id)
173            .map(|adjacent| {
174                adjacent
175                    .iter()
176                    .filter(|(_, edge)| edge.ast.is_some())
177                    .map(|(to, _)| *to)
178                    .collect()
179            })
180            .unwrap_or_default()
181    }
182
183    fn parents(&self, n_id: NodeId) -> Vec<NodeId> {
184        self.edges
185            .iter()
186            .filter(|(_, adjacent)| adjacent.contains_key(&n_id))
187            .map(|(from, _)| *from)
188            .collect()
189    }
190
191    fn ast_parents(&self, n_id: NodeId) -> Vec<NodeId> {
192        self.edges
193            .iter()
194            .filter(|(_, adjacent)| adjacent.get(&n_id).is_some_and(|edge| edge.ast.is_some()))
195            .map(|(from, _)| *from)
196            .collect()
197    }
198
199    fn cfg_parents(&self, n_id: NodeId) -> Vec<NodeId> {
200        self.edges
201            .iter()
202            .filter(|(_, adjacent)| adjacent.get(&n_id).is_some_and(|edge| edge.cfg.is_some()))
203            .map(|(from, _)| *from)
204            .collect()
205    }
206
207    fn cfg_children(&self, n_id: NodeId) -> Vec<NodeId> {
208        self.edges
209            .get(&n_id)
210            .map(|adjacent| {
211                adjacent
212                    .iter()
213                    .filter(|(_, edge)| edge.cfg.is_some())
214                    .map(|(to, _)| *to)
215                    .collect()
216            })
217            .unwrap_or_default()
218    }
219
220    fn label_type(&self, n_id: NodeId) -> Option<&str> {
221        self.nodes.get(&n_id).map(SyntaxNode::label_type)
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use super::SyntaxGraph;
228    use crate::syntax::{SyntaxEdge, SyntaxKind, SyntaxNode};
229    use crate::{Ast, Cfg, CodeGraph, NodeId};
230    use alloc::borrow::ToOwned;
231
232    #[test]
233    fn stores_nodes_and_edges_by_id() {
234        let mut graph = SyntaxGraph::new();
235        graph.add_node(
236            NodeId(1),
237            SyntaxNode::MissingNode {
238                node_type: "stream".to_owned(),
239            },
240        );
241        graph.add_ast_edge(NodeId(1), NodeId(2));
242
243        assert_eq!(
244            graph.nodes.get(&NodeId(1)),
245            Some(&SyntaxNode::MissingNode {
246                node_type: "stream".to_owned()
247            })
248        );
249        assert_eq!(
250            graph
251                .edges
252                .get(&NodeId(1))
253                .and_then(|adjacent| adjacent.get(&NodeId(2))),
254            Some(&SyntaxEdge {
255                ast: Some(Ast),
256                cfg: None
257            })
258        );
259    }
260
261    #[test]
262    fn add_node_indexes_ids_by_syntax_type_in_insertion_order() {
263        let mut graph = SyntaxGraph::new();
264        graph.add_node(NodeId(9), SyntaxNode::Break);
265        graph.add_node(NodeId(1), SyntaxNode::File);
266        graph.add_node(NodeId(2), SyntaxNode::Break);
267
268        assert_eq!(
269            graph.nodes_by_type.get(&SyntaxKind::Break),
270            Some(&alloc::vec![NodeId(9), NodeId(2)])
271        );
272        assert_eq!(
273            graph.nodes_by_type.get(&SyntaxKind::File),
274            Some(&alloc::vec![NodeId(1)])
275        );
276        assert_eq!(graph.nodes_by_type.get(&SyntaxKind::If), None);
277    }
278
279    #[test]
280    fn ast_edge_merges_into_existing_edge() {
281        let mut graph = SyntaxGraph::new();
282        graph.add_ast_edge(NodeId(1), NodeId(2));
283        graph.add_ast_edge(NodeId(1), NodeId(2));
284
285        assert_eq!(
286            graph
287                .edges
288                .get(&NodeId(1))
289                .map(alloc::collections::BTreeMap::len),
290            Some(1)
291        );
292    }
293
294    #[test]
295    fn ast_children_skip_edges_without_the_ast_mark() {
296        let mut graph = SyntaxGraph::new();
297        graph.add_ast_edge(NodeId(1), NodeId(2));
298        graph.edges.entry(NodeId(1)).or_default().insert(
299            NodeId(3),
300            SyntaxEdge {
301                ast: None,
302                cfg: None,
303            },
304        );
305
306        assert_eq!(graph.ast_children(NodeId(1)), [NodeId(2)]);
307    }
308
309    #[test]
310    fn cfg_edge_merges_onto_the_ast_edge() {
311        let mut graph = SyntaxGraph::new();
312        graph.add_ast_edge(NodeId(1), NodeId(2));
313        graph.add_cfg_edge(NodeId(1), NodeId(2));
314
315        assert_eq!(
316            graph
317                .edges
318                .get(&NodeId(1))
319                .and_then(|adjacent| adjacent.get(&NodeId(2))),
320            Some(&SyntaxEdge {
321                ast: Some(Ast),
322                cfg: Some(Cfg)
323            })
324        );
325    }
326
327    #[test]
328    fn remove_edge_drops_the_edge_and_prunes_the_empty_entry() {
329        let mut graph = SyntaxGraph::new();
330        graph.add_ast_edge(NodeId(1), NodeId(2));
331        graph.add_ast_edge(NodeId(1), NodeId(3));
332        graph.add_cfg_edge(NodeId(1), NodeId(2));
333
334        graph.remove_edge(NodeId(1), NodeId(2));
335
336        assert_eq!(graph.children(NodeId(1)), [NodeId(3)]);
337        assert_eq!(graph.parents(NodeId(2)), []);
338
339        graph.remove_edge(NodeId(1), NodeId(3));
340
341        assert!(!graph.edges.contains_key(&NodeId(1)));
342    }
343
344    #[test]
345    fn finalize_symbol_index_puts_the_most_recent_definition_first() {
346        let mut graph = SyntaxGraph::new();
347        graph
348            .symbol_index
349            .entry("x".to_owned())
350            .or_default()
351            .entry(NodeId(1))
352            .or_default()
353            .extend([NodeId(2), NodeId(5), NodeId(9)]);
354
355        graph.finalize_symbol_index();
356
357        assert_eq!(
358            graph
359                .symbol_index
360                .get("x")
361                .and_then(|scopes| scopes.get(&NodeId(1))),
362            Some(&alloc::vec![NodeId(9), NodeId(5), NodeId(2)])
363        );
364    }
365
366    #[test]
367    fn remove_edge_ignores_a_missing_edge() {
368        let mut graph = SyntaxGraph::new();
369        graph.add_ast_edge(NodeId(1), NodeId(2));
370
371        graph.remove_edge(NodeId(7), NodeId(8));
372        graph.remove_edge(NodeId(1), NodeId(9));
373
374        assert_eq!(graph.children(NodeId(1)), [NodeId(2)]);
375    }
376
377    #[test]
378    fn cfg_children_skip_edges_without_the_cfg_mark() {
379        let mut graph = SyntaxGraph::new();
380        graph.add_ast_edge(NodeId(1), NodeId(2));
381        graph.add_ast_edge(NodeId(1), NodeId(3));
382        graph.add_cfg_edge(NodeId(1), NodeId(3));
383
384        assert_eq!(graph.cfg_children(NodeId(1)), [NodeId(3)]);
385        assert_eq!(graph.cfg_children(NodeId(9)), []);
386    }
387
388    #[test]
389    fn register_call_groups_invocations_by_name() {
390        let mut graph = SyntaxGraph::new();
391        graph.register_call("Render", NodeId(5));
392        graph.register_call("Render", NodeId(9));
393        assert_eq!(
394            graph.call_index.get("Render").map(Vec::as_slice),
395            Some([NodeId(5), NodeId(9)].as_slice())
396        );
397    }
398}