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/// The AST graph: every tree-sitter node plus parent-to-child adjacency.
12#[derive(Clone, PartialEq, Eq, Debug, Default)]
13#[non_exhaustive]
14pub struct AstGraph {
15    /// Every node, keyed by id.
16    pub nodes: BTreeMap<NodeId, AstNode>,
17    /// Adjacency: parent id -> child id -> the edge connecting them.
18    pub edges: BTreeMap<NodeId, BTreeMap<NodeId, AstEdge>>,
19}
20
21impl AstGraph {
22    /// An empty AST graph.
23    #[must_use]
24    pub fn new() -> Self {
25        Self::default()
26    }
27
28    /// Inserts `node` under `id`, replacing any node already there.
29    pub fn add_node(&mut self, id: NodeId, node: AstNode) {
30        self.nodes.insert(id, node);
31    }
32
33    /// Connects parent `from` to child `to` at the grammar's child `index`.
34    pub fn add_edge(&mut self, from: NodeId, to: NodeId, index: u32) {
35        // The adjacency map doubles as the child set, so a repeated edge never
36        // reaches the ordered `children` list.
37        let is_new_child = self
38            .edges
39            .entry(from)
40            .or_default()
41            .insert(to, AstEdge { kind: Ast, index })
42            .is_none();
43        if is_new_child {
44            if let Some(node) = self.nodes.get_mut(&from) {
45                node.children.push(to);
46            }
47        }
48    }
49
50    /// Names `child` as `parent`'s grammar field `name`, keeping the first write.
51    pub fn set_field(&mut self, parent: NodeId, name: String, child: NodeId) {
52        if let Some(node) = self.nodes.get_mut(&parent) {
53            node.fields.entry(name).or_insert(child);
54        }
55    }
56}
57
58impl CodeGraph for AstGraph {
59    fn children(&self, n_id: NodeId) -> Vec<NodeId> {
60        self.ast_children(n_id)
61    }
62
63    fn ast_children(&self, n_id: NodeId) -> Vec<NodeId> {
64        self.edges
65            .get(&n_id)
66            .map(|adjacent| adjacent.keys().copied().collect())
67            .unwrap_or_default()
68    }
69
70    fn parents(&self, n_id: NodeId) -> Vec<NodeId> {
71        self.ast_parents(n_id)
72    }
73
74    fn ast_parents(&self, n_id: NodeId) -> Vec<NodeId> {
75        self.edges
76            .iter()
77            .filter(|(_, adjacent)| adjacent.contains_key(&n_id))
78            .map(|(from, _)| *from)
79            .collect()
80    }
81
82    fn cfg_children(&self, _n_id: NodeId) -> Vec<NodeId> {
83        Vec::new()
84    }
85
86    fn cfg_parents(&self, _n_id: NodeId) -> Vec<NodeId> {
87        Vec::new()
88    }
89
90    fn label_type(&self, n_id: NodeId) -> Option<&str> {
91        self.nodes.get(&n_id).map(|node| node.kind.as_str())
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::AstGraph;
98    use crate::ast::AstNode;
99    use crate::{Ast, CodeGraph, NodeId};
100    use alloc::borrow::ToOwned;
101    use alloc::vec::Vec;
102
103    #[test]
104    fn builds_parent_child_with_field_and_text() {
105        let mut graph = AstGraph::new();
106
107        graph.add_node(
108            NodeId(1),
109            AstNode::new(1, 1, "method_declaration".to_owned()),
110        );
111        let mut name = AstNode::new(1, 8, "identifier".to_owned());
112        name.text = Some("foo".to_owned());
113        graph.add_node(NodeId(2), name);
114
115        graph.add_edge(NodeId(1), NodeId(2), 0);
116        graph.set_field(NodeId(1), "name".to_owned(), NodeId(2));
117
118        assert_eq!(graph.nodes.len(), 2);
119        assert_eq!(
120            graph.nodes.get(&NodeId(2)).and_then(|n| n.text.as_deref()),
121            Some("foo")
122        );
123        assert_eq!(
124            graph
125                .nodes
126                .get(&NodeId(1))
127                .and_then(|n| n.fields.get("name")),
128            Some(&NodeId(2))
129        );
130        let edge = graph
131            .edges
132            .get(&NodeId(1))
133            .and_then(|m| m.get(&NodeId(2)))
134            .copied();
135        assert_eq!(edge.map(|e| e.index), Some(0));
136        assert_eq!(edge.map(|e| e.kind), Some(Ast));
137    }
138
139    #[test]
140    fn set_field_on_unknown_parent_is_noop() {
141        let mut graph = AstGraph::new();
142        graph.set_field(NodeId(99), "name".to_owned(), NodeId(1));
143        assert!(graph.nodes.is_empty());
144    }
145
146    #[test]
147    fn set_field_keeps_first_write() {
148        let mut graph = AstGraph::new();
149        graph.add_node(NodeId(1), AstNode::new(1, 1, "x".to_owned()));
150        graph.set_field(NodeId(1), "name".to_owned(), NodeId(2));
151        graph.set_field(NodeId(1), "name".to_owned(), NodeId(3));
152        assert_eq!(
153            graph
154                .nodes
155                .get(&NodeId(1))
156                .and_then(|n| n.fields.get("name")),
157            Some(&NodeId(2))
158        );
159    }
160
161    #[test]
162    fn iteration_order_is_node_id_order() {
163        let mut graph = AstGraph::new();
164        graph.add_node(NodeId(2), AstNode::new(1, 1, "b".to_owned()));
165        graph.add_node(NodeId(10), AstNode::new(1, 1, "c".to_owned()));
166        graph.add_node(NodeId(1), AstNode::new(1, 1, "a".to_owned()));
167
168        let ids: Vec<u64> = graph.nodes.keys().map(|n| n.0).collect();
169        assert_eq!(ids, [1, 2, 10]);
170    }
171
172    #[test]
173    fn ast_children_are_all_edges_in_id_order() {
174        let mut graph = AstGraph::new();
175        graph.add_node(NodeId(1), AstNode::new(1, 1, "root".to_owned()));
176        graph.add_edge(NodeId(1), NodeId(3), 1);
177        graph.add_edge(NodeId(1), NodeId(2), 0);
178
179        assert_eq!(graph.ast_children(NodeId(1)), [NodeId(2), NodeId(3)]);
180        assert_eq!(graph.ast_children(NodeId(9)), []);
181        assert_eq!(graph.label_type(NodeId(1)), Some("root"));
182        assert_eq!(graph.label_type(NodeId(9)), None);
183    }
184
185    #[test]
186    fn add_edge_records_children_in_insertion_order() {
187        let mut graph = AstGraph::new();
188        graph.add_node(NodeId(1), AstNode::new(1, 1, "root".to_owned()));
189        graph.add_edge(NodeId(1), NodeId(3), 0);
190        graph.add_edge(NodeId(1), NodeId(2), 1);
191        graph.add_edge(NodeId(1), NodeId(3), 0);
192
193        assert_eq!(
194            graph.nodes.get(&NodeId(1)).map(|n| n.children.as_slice()),
195            Some([NodeId(3), NodeId(2)].as_slice())
196        );
197    }
198}
fluidattacks_blends_domain/ast/graph.rs

fluidattacks_blends_domain/ast/
graph.rs
