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

fluidattacks_blends_domain/ast/
graph.rs
