fluidattacks_blends_domain/ast/
graph.rs1use 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}