Skip to main content

fluidattacks_blends_domain/ast/
graph.rs

1//! AST graph: node table plus parent → child adjacency.
2
3use alloc::collections::BTreeMap;
4use alloc::string::String;
5use alloc::vec::Vec;
6
7use crate::ast::AstEdge;
8use crate::ast::AstNode;
9use crate::{Ast, CodeGraph, NodeId};
10
11#[derive(Clone, PartialEq, Eq, Debug, Default)]
12pub struct AstGraph {
13    pub nodes: BTreeMap<NodeId, AstNode>,
14    pub edges: BTreeMap<NodeId, BTreeMap<NodeId, AstEdge>>,
15}
16
17impl AstGraph {
18    #[must_use]
19    pub fn new() -> Self {
20        Self::default()
21    }
22
23    pub fn add_node(&mut self, id: NodeId, node: AstNode) {
24        self.nodes.insert(id, node);
25    }
26
27    pub fn add_edge(&mut self, from: NodeId, to: NodeId, index: u32) {
28        self.edges
29            .entry(from)
30            .or_default()
31            .insert(to, AstEdge { kind: Ast, index });
32    }
33
34    pub fn set_field(&mut self, parent: NodeId, name: String, child: NodeId) {
35        if let Some(node) = self.nodes.get_mut(&parent) {
36            node.fields.entry(name).or_insert(child);
37        }
38    }
39}
40
41impl CodeGraph for AstGraph {
42    fn children(&self, n_id: NodeId) -> Vec<NodeId> {
43        self.ast_children(n_id)
44    }
45
46    fn ast_children(&self, n_id: NodeId) -> Vec<NodeId> {
47        self.edges
48            .get(&n_id)
49            .map(|adjacent| adjacent.keys().copied().collect())
50            .unwrap_or_default()
51    }
52
53    fn parents(&self, n_id: NodeId) -> Vec<NodeId> {
54        self.ast_parents(n_id)
55    }
56
57    fn ast_parents(&self, n_id: NodeId) -> Vec<NodeId> {
58        self.edges
59            .iter()
60            .filter(|(_, adjacent)| adjacent.contains_key(&n_id))
61            .map(|(from, _)| *from)
62            .collect()
63    }
64
65    fn cfg_children(&self, _n_id: NodeId) -> Vec<NodeId> {
66        Vec::new()
67    }
68
69    fn cfg_parents(&self, _n_id: NodeId) -> Vec<NodeId> {
70        Vec::new()
71    }
72
73    fn label_type(&self, n_id: NodeId) -> Option<&str> {
74        self.nodes.get(&n_id).map(|node| node.kind.as_str())
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::AstGraph;
81    use crate::ast::AstNode;
82    use crate::{Ast, CodeGraph, NodeId};
83    use alloc::borrow::ToOwned;
84    use alloc::vec::Vec;
85
86    #[test]
87    fn builds_parent_child_with_field_and_text() {
88        let mut graph = AstGraph::new();
89
90        graph.add_node(
91            NodeId(1),
92            AstNode::new(1, 1, "method_declaration".to_owned()),
93        );
94        let mut name = AstNode::new(1, 8, "identifier".to_owned());
95        name.text = Some("foo".to_owned());
96        graph.add_node(NodeId(2), name);
97
98        graph.add_edge(NodeId(1), NodeId(2), 0);
99        graph.set_field(NodeId(1), "name".to_owned(), NodeId(2));
100
101        assert_eq!(graph.nodes.len(), 2);
102        assert_eq!(
103            graph.nodes.get(&NodeId(2)).and_then(|n| n.text.as_deref()),
104            Some("foo")
105        );
106        assert_eq!(
107            graph
108                .nodes
109                .get(&NodeId(1))
110                .and_then(|n| n.fields.get("name")),
111            Some(&NodeId(2))
112        );
113        let edge = graph
114            .edges
115            .get(&NodeId(1))
116            .and_then(|m| m.get(&NodeId(2)))
117            .copied();
118        assert_eq!(edge.map(|e| e.index), Some(0));
119        assert_eq!(edge.map(|e| e.kind), Some(Ast));
120    }
121
122    #[test]
123    fn set_field_on_unknown_parent_is_noop() {
124        let mut graph = AstGraph::new();
125        graph.set_field(NodeId(99), "name".to_owned(), NodeId(1));
126        assert!(graph.nodes.is_empty());
127    }
128
129    #[test]
130    fn set_field_keeps_first_write() {
131        let mut graph = AstGraph::new();
132        graph.add_node(NodeId(1), AstNode::new(1, 1, "x".to_owned()));
133        graph.set_field(NodeId(1), "name".to_owned(), NodeId(2));
134        graph.set_field(NodeId(1), "name".to_owned(), NodeId(3));
135        assert_eq!(
136            graph
137                .nodes
138                .get(&NodeId(1))
139                .and_then(|n| n.fields.get("name")),
140            Some(&NodeId(2))
141        );
142    }
143
144    #[test]
145    fn iteration_order_is_node_id_order() {
146        let mut graph = AstGraph::new();
147        graph.add_node(NodeId(2), AstNode::new(1, 1, "b".to_owned()));
148        graph.add_node(NodeId(10), AstNode::new(1, 1, "c".to_owned()));
149        graph.add_node(NodeId(1), AstNode::new(1, 1, "a".to_owned()));
150
151        let ids: Vec<u64> = graph.nodes.keys().map(|n| n.0).collect();
152        assert_eq!(ids, [1, 2, 10]);
153    }
154
155    #[test]
156    fn ast_children_are_all_edges_in_id_order() {
157        let mut graph = AstGraph::new();
158        graph.add_node(NodeId(1), AstNode::new(1, 1, "root".to_owned()));
159        graph.add_edge(NodeId(1), NodeId(3), 1);
160        graph.add_edge(NodeId(1), NodeId(2), 0);
161
162        assert_eq!(graph.ast_children(NodeId(1)), [NodeId(2), NodeId(3)]);
163        assert_eq!(graph.ast_children(NodeId(9)), []);
164        assert_eq!(graph.label_type(NodeId(1)), Some("root"));
165        assert_eq!(graph.label_type(NodeId(9)), None);
166    }
167}