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

fluidattacks_blends_domain/ast/
graph.rs
