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    /// Node id -> enclosing (if id, is true branch) pairs, outermost first.
70    pub branch_index: BTreeMap<NodeId, Vec<(NodeId, bool)>>,
71    /// Callee short name -> invocations of it, used to reach a method's call sites.
72    pub call_index: BTreeMap<String, Vec<NodeId>>,
73    /// Syntax kind -> ids of every node of that kind, in insertion order.
74    pub nodes_by_type: BTreeMap<SyntaxKind, Vec<NodeId>>,
75}
76
77impl SyntaxGraph {
78    /// An empty syntax graph.
79    #[must_use]
80    pub fn new() -> Self {
81        Self::default()
82    }
83
84    /// Inserts `node` under `id`, indexing it by its syntax kind.
85    pub fn add_node(&mut self, id: NodeId, node: SyntaxNode) {
86        self.nodes_by_type
87            .entry(SyntaxKind::from(&node))
88            .or_default()
89            .push(id);
90        self.nodes.insert(id, node);
91    }
92
93    /// Records a sanitization of `symbol` as seen within `scope`.
94    pub fn register_sanitization(&mut self, symbol: &str, scope: NodeId, event: SanitizationEvent) {
95        self.sanitization_index
96            .entry(String::from(symbol))
97            .or_default()
98            .entry(scope)
99            .or_default()
100            .push(event);
101    }
102
103    /// Records that `invocation_id` calls a callee named `name`.
104    pub fn register_call(&mut self, name: &str, invocation_id: NodeId) {
105        self.call_index
106            .entry(String::from(name))
107            .or_default()
108            .push(invocation_id);
109    }
110
111    /// Records a use of `symbol` as seen within `scope`.
112    pub fn register_usage(&mut self, symbol: &str, scope: NodeId, event: UsageEvent) {
113        self.usage_index
114            .entry(String::from(symbol))
115            .or_default()
116            .entry(scope)
117            .or_default()
118            .push(event);
119    }
120
121    /// Connects `from` to `to` with an AST edge, merging into any existing edge.
122    pub fn add_ast_edge(&mut self, from: NodeId, to: NodeId) {
123        self.edges
124            .entry(from)
125            .or_default()
126            .entry(to)
127            .or_default()
128            .ast = Some(Ast);
129    }
130
131    /// Connects `from` to `to` with a CFG edge, merging into any existing edge.
132    pub fn add_cfg_edge(&mut self, from: NodeId, to: NodeId) {
133        self.edges
134            .entry(from)
135            .or_default()
136            .entry(to)
137            .or_default()
138            .cfg = Some(Cfg);
139    }
140
141    /// Reverses each scope's definition list so lookups see the most recent first.
142    ///
143    /// Call once after the graph is fully built; definitions are appended in
144    /// parse order, which is the opposite of the shadowing order callers need.
145    pub fn finalize_symbol_index(&mut self) {
146        for scope_definitions in self.symbol_index.values_mut() {
147            for definition_list in scope_definitions.values_mut() {
148                definition_list.reverse();
149            }
150        }
151    }
152
153    /// Removes the edge from `from` to `to`, if one exists.
154    pub fn remove_edge(&mut self, from: NodeId, to: NodeId) {
155        if let Some(adjacent) = self.edges.get_mut(&from) {
156            adjacent.remove(&to);
157            if adjacent.is_empty() {
158                self.edges.remove(&from);
159            }
160        }
161    }
162}
163
164impl CodeGraph for SyntaxGraph {
165    fn children(&self, n_id: NodeId) -> Vec<NodeId> {
166        self.edges
167            .get(&n_id)
168            .map(|adjacent| adjacent.keys().copied().collect())
169            .unwrap_or_default()
170    }
171
172    fn ast_children(&self, n_id: NodeId) -> Vec<NodeId> {
173        self.edges
174            .get(&n_id)
175            .map(|adjacent| {
176                adjacent
177                    .iter()
178                    .filter(|(_, edge)| edge.ast.is_some())
179                    .map(|(to, _)| *to)
180                    .collect()
181            })
182            .unwrap_or_default()
183    }
184
185    fn parents(&self, n_id: NodeId) -> Vec<NodeId> {
186        self.edges
187            .iter()
188            .filter(|(_, adjacent)| adjacent.contains_key(&n_id))
189            .map(|(from, _)| *from)
190            .collect()
191    }
192
193    fn ast_parents(&self, n_id: NodeId) -> Vec<NodeId> {
194        self.edges
195            .iter()
196            .filter(|(_, adjacent)| adjacent.get(&n_id).is_some_and(|edge| edge.ast.is_some()))
197            .map(|(from, _)| *from)
198            .collect()
199    }
200
201    fn cfg_parents(&self, n_id: NodeId) -> Vec<NodeId> {
202        self.edges
203            .iter()
204            .filter(|(_, adjacent)| adjacent.get(&n_id).is_some_and(|edge| edge.cfg.is_some()))
205            .map(|(from, _)| *from)
206            .collect()
207    }
208
209    fn cfg_children(&self, n_id: NodeId) -> Vec<NodeId> {
210        self.edges
211            .get(&n_id)
212            .map(|adjacent| {
213                adjacent
214                    .iter()
215                    .filter(|(_, edge)| edge.cfg.is_some())
216                    .map(|(to, _)| *to)
217                    .collect()
218            })
219            .unwrap_or_default()
220    }
221
222    fn label_type(&self, n_id: NodeId) -> Option<&str> {
223        self.nodes.get(&n_id).map(SyntaxNode::label_type)
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use super::SyntaxGraph;
230    use crate::syntax::{SyntaxEdge, SyntaxKind, SyntaxNode};
231    use crate::{Ast, Cfg, CodeGraph, NodeId};
232    use alloc::borrow::ToOwned;
233
234    #[test]
235    fn stores_nodes_and_edges_by_id() {
236        let mut graph = SyntaxGraph::new();
237        graph.add_node(
238            NodeId(1),
239            SyntaxNode::MissingNode {
240                node_type: "stream".to_owned(),
241            },
242        );
243        graph.add_ast_edge(NodeId(1), NodeId(2));
244
245        assert_eq!(
246            graph.nodes.get(&NodeId(1)),
247            Some(&SyntaxNode::MissingNode {
248                node_type: "stream".to_owned()
249            })
250        );
251        assert_eq!(
252            graph
253                .edges
254                .get(&NodeId(1))
255                .and_then(|adjacent| adjacent.get(&NodeId(2))),
256            Some(&SyntaxEdge {
257                ast: Some(Ast),
258                cfg: None
259            })
260        );
261    }
262
263    #[test]
264    fn add_node_indexes_ids_by_syntax_type_in_insertion_order() {
265        let mut graph = SyntaxGraph::new();
266        graph.add_node(NodeId(9), SyntaxNode::Break);
267        graph.add_node(NodeId(1), SyntaxNode::File);
268        graph.add_node(NodeId(2), SyntaxNode::Break);
269
270        assert_eq!(
271            graph.nodes_by_type.get(&SyntaxKind::Break),
272            Some(&alloc::vec![NodeId(9), NodeId(2)])
273        );
274        assert_eq!(
275            graph.nodes_by_type.get(&SyntaxKind::File),
276            Some(&alloc::vec![NodeId(1)])
277        );
278        assert_eq!(graph.nodes_by_type.get(&SyntaxKind::If), None);
279    }
280
281    #[test]
282    fn ast_edge_merges_into_existing_edge() {
283        let mut graph = SyntaxGraph::new();
284        graph.add_ast_edge(NodeId(1), NodeId(2));
285        graph.add_ast_edge(NodeId(1), NodeId(2));
286
287        assert_eq!(
288            graph
289                .edges
290                .get(&NodeId(1))
291                .map(alloc::collections::BTreeMap::len),
292            Some(1)
293        );
294    }
295
296    #[test]
297    fn ast_children_skip_edges_without_the_ast_mark() {
298        let mut graph = SyntaxGraph::new();
299        graph.add_ast_edge(NodeId(1), NodeId(2));
300        graph.edges.entry(NodeId(1)).or_default().insert(
301            NodeId(3),
302            SyntaxEdge {
303                ast: None,
304                cfg: None,
305            },
306        );
307
308        assert_eq!(graph.ast_children(NodeId(1)), [NodeId(2)]);
309    }
310
311    #[test]
312    fn cfg_edge_merges_onto_the_ast_edge() {
313        let mut graph = SyntaxGraph::new();
314        graph.add_ast_edge(NodeId(1), NodeId(2));
315        graph.add_cfg_edge(NodeId(1), NodeId(2));
316
317        assert_eq!(
318            graph
319                .edges
320                .get(&NodeId(1))
321                .and_then(|adjacent| adjacent.get(&NodeId(2))),
322            Some(&SyntaxEdge {
323                ast: Some(Ast),
324                cfg: Some(Cfg)
325            })
326        );
327    }
328
329    #[test]
330    fn remove_edge_drops_the_edge_and_prunes_the_empty_entry() {
331        let mut graph = SyntaxGraph::new();
332        graph.add_ast_edge(NodeId(1), NodeId(2));
333        graph.add_ast_edge(NodeId(1), NodeId(3));
334        graph.add_cfg_edge(NodeId(1), NodeId(2));
335
336        graph.remove_edge(NodeId(1), NodeId(2));
337
338        assert_eq!(graph.children(NodeId(1)), [NodeId(3)]);
339        assert_eq!(graph.parents(NodeId(2)), []);
340
341        graph.remove_edge(NodeId(1), NodeId(3));
342
343        assert!(!graph.edges.contains_key(&NodeId(1)));
344    }
345
346    #[test]
347    fn finalize_symbol_index_puts_the_most_recent_definition_first() {
348        let mut graph = SyntaxGraph::new();
349        graph
350            .symbol_index
351            .entry("x".to_owned())
352            .or_default()
353            .entry(NodeId(1))
354            .or_default()
355            .extend([NodeId(2), NodeId(5), NodeId(9)]);
356
357        graph.finalize_symbol_index();
358
359        assert_eq!(
360            graph
361                .symbol_index
362                .get("x")
363                .and_then(|scopes| scopes.get(&NodeId(1))),
364            Some(&alloc::vec![NodeId(9), NodeId(5), NodeId(2)])
365        );
366    }
367
368    #[test]
369    fn remove_edge_ignores_a_missing_edge() {
370        let mut graph = SyntaxGraph::new();
371        graph.add_ast_edge(NodeId(1), NodeId(2));
372
373        graph.remove_edge(NodeId(7), NodeId(8));
374        graph.remove_edge(NodeId(1), NodeId(9));
375
376        assert_eq!(graph.children(NodeId(1)), [NodeId(2)]);
377    }
378
379    #[test]
380    fn cfg_children_skip_edges_without_the_cfg_mark() {
381        let mut graph = SyntaxGraph::new();
382        graph.add_ast_edge(NodeId(1), NodeId(2));
383        graph.add_ast_edge(NodeId(1), NodeId(3));
384        graph.add_cfg_edge(NodeId(1), NodeId(3));
385
386        assert_eq!(graph.cfg_children(NodeId(1)), [NodeId(3)]);
387        assert_eq!(graph.cfg_children(NodeId(9)), []);
388    }
389
390    #[test]
391    fn register_call_groups_invocations_by_name() {
392        let mut graph = SyntaxGraph::new();
393        graph.register_call("Render", NodeId(5));
394        graph.register_call("Render", NodeId(9));
395        assert_eq!(
396            graph.call_index.get("Render").map(Vec::as_slice),
397            Some([NodeId(5), NodeId(9)].as_slice())
398        );
399    }
400}