1use std::collections::HashMap;
2
3use camino::Utf8Path;
4
5use crate::edge::{Edge, EdgeKind};
6use crate::selector::SymbolSelector;
7
8#[derive(Debug, Clone)]
9pub struct SymbolOutline {
10 pub name: String,
11 pub kind: String,
12 pub signature: Option<String>,
13 pub range: (u64, u64),
14 pub children: Vec<SymbolOutline>,
15}
16
17#[derive(Debug)]
18pub struct Graph {
19 edges: Vec<Edge>,
20 by_from: HashMap<String, Vec<usize>>,
21 by_to: HashMap<String, Vec<usize>>,
22 by_kind: HashMap<EdgeKind, Vec<usize>>,
23}
24
25impl Graph {
26 #[must_use]
27 pub fn new(edges: Vec<Edge>) -> Self {
28 let mut by_from: HashMap<String, Vec<usize>> = HashMap::new();
29 let mut by_to: HashMap<String, Vec<usize>> = HashMap::new();
30 let mut by_kind: HashMap<EdgeKind, Vec<usize>> = HashMap::new();
31
32 for (i, edge) in edges.iter().enumerate() {
33 by_from
34 .entry(edge.from.as_str().to_string())
35 .or_default()
36 .push(i);
37 by_to
38 .entry(edge.to.as_str().to_string())
39 .or_default()
40 .push(i);
41 by_kind.entry(edge.kind).or_default().push(i);
42 }
43
44 Self {
45 edges,
46 by_from,
47 by_to,
48 by_kind,
49 }
50 }
51
52 #[must_use]
53 pub fn edges(&self) -> &[Edge] {
54 &self.edges
55 }
56
57 #[must_use]
58 pub fn count_by_kind(&self, kind: EdgeKind) -> usize {
59 self.by_kind.get(&kind).map_or(0, |v| v.len())
60 }
61
62 pub fn find_definition(&self, name: &str, file: Option<&Utf8Path>) -> Vec<&Edge> {
63 let def_indexes = self.by_kind.get(&EdgeKind::Defines);
64 let Some(def_indexes) = def_indexes else {
65 return Vec::new();
66 };
67
68 def_indexes
69 .iter()
70 .map(|&i| &self.edges[i])
71 .filter(|e| {
72 let id_str = e.to.as_str();
73 let contains_name = id_str.contains(&format!("::{name}::"));
74 let matches_file = file.is_none_or(|f| {
75 let f_str = f.as_str();
76 id_str.starts_with(f_str)
77 || f_str.ends_with(id_str.split("::").next().unwrap_or(""))
78 });
79 contains_name && matches_file
80 })
81 .collect()
82 }
83
84 pub fn find_references<'a>(&'a self, sel: &SymbolSelector) -> Vec<&'a Edge> {
85 let target_ids = self.resolve_selector(sel);
86 self.edges_matching(target_ids, EdgeKind::References)
87 }
88
89 pub fn callers<'a>(&'a self, sel: &SymbolSelector) -> Vec<&'a Edge> {
90 let target_ids = self.resolve_selector(sel);
91 self.edges_matching(target_ids, EdgeKind::Calls)
92 }
93
94 pub fn callees<'a>(&'a self, sel: &SymbolSelector) -> Vec<&'a Edge> {
95 let source_ids = self.resolve_selector(sel);
96 self.edges_matching_from(source_ids, EdgeKind::Calls)
97 }
98
99 pub fn imports_of<'a>(&'a self, file: &Utf8Path) -> Vec<&'a Edge> {
100 let import_indexes = self.by_kind.get(&EdgeKind::Imports);
101 let Some(import_indexes) = import_indexes else {
102 return Vec::new();
103 };
104 let file_str = file.as_str();
105 import_indexes
106 .iter()
107 .map(|&i| &self.edges[i])
108 .filter(|e| {
109 let id_str = e.from.as_str();
110 id_str.starts_with(file_str)
111 || file_str.ends_with(id_str.split("::").next().unwrap_or(""))
112 })
113 .collect()
114 }
115
116 pub fn imported_by<'a>(&'a self, sel: &SymbolSelector) -> Vec<&'a Edge> {
117 let target_ids = self.resolve_selector(sel);
118 self.edges_matching_from(target_ids, EdgeKind::ImportedBy)
119 }
120
121 pub fn outline(&self, file: &Utf8Path) -> Vec<SymbolOutline> {
122 let def_indexes = self.by_kind.get(&EdgeKind::Defines);
123 let Some(def_indexes) = def_indexes else {
124 return Vec::new();
125 };
126
127 let file_str = file.as_str();
128 def_indexes
129 .iter()
130 .map(|&i| &self.edges[i])
131 .filter(|e| {
132 let id_str = e.from.as_str();
133 id_str.starts_with(file_str)
134 || file_str.ends_with(id_str.split("::").next().unwrap_or(""))
135 })
136 .map(|e| {
137 let id_str = e.from.as_str();
138 let name = id_str.split("::").nth(1).unwrap_or("?");
139 SymbolOutline {
140 name: name.to_string(),
141 kind: "symbol".to_string(),
142 signature: None,
143 range: (0, 0),
144 children: Vec::new(),
145 }
146 })
147 .collect()
148 }
149
150 fn resolve_selector(&self, sel: &SymbolSelector) -> Vec<String> {
151 match sel {
152 SymbolSelector::ById(id) => vec![id.as_str().to_string()],
153 SymbolSelector::ByName { file, name } => {
154 let prefix = format!("{file}::{name}::");
155 self.edges
156 .iter()
157 .filter(|e| e.to.as_str().starts_with(&prefix))
158 .map(|e| e.to.as_str().to_string())
159 .collect()
160 }
161 SymbolSelector::Qualified(qn) => self
162 .edges
163 .iter()
164 .filter(|e| e.to.as_str().contains(qn.as_str()))
165 .map(|e| e.to.as_str().to_string())
166 .collect(),
167 }
168 }
169
170 fn edges_matching(&self, target_ids: Vec<String>, kind: EdgeKind) -> Vec<&Edge> {
171 let mut results = Vec::new();
172 for tid in &target_ids {
173 if let Some(indexes) = self.by_to.get(tid) {
174 for &i in indexes {
175 let edge = &self.edges[i];
176 if edge.kind == kind {
177 results.push(edge);
178 }
179 }
180 }
181 }
182 results
183 }
184
185 fn edges_matching_from(&self, source_ids: Vec<String>, kind: EdgeKind) -> Vec<&Edge> {
186 let mut results = Vec::new();
187 for sid in &source_ids {
188 if let Some(indexes) = self.by_from.get(sid) {
189 for &i in indexes {
190 let edge = &self.edges[i];
191 if edge.kind == kind {
192 results.push(edge);
193 }
194 }
195 }
196 }
197 results
198 }
199}
200
201#[cfg(test)]
202#[allow(clippy::unwrap_used, clippy::expect_used)]
203mod tests {
204 use super::*;
205 use crate::edge::{Confidence, Edge, EdgeKind};
206 use argyph_parse::SymbolId;
207
208 #[test]
209 fn new_graph_is_queryable() {
210 let edges = vec![Edge {
211 from: SymbolId::new(&camino::Utf8PathBuf::from("src/lib.rs"), "fn_a", 0),
212 to: SymbolId::new(&camino::Utf8PathBuf::from("src/lib.rs"), "fn_a", 0),
213 kind: EdgeKind::Defines,
214 confidence: Confidence::Resolved,
215 }];
216 let graph = Graph::new(edges);
217 assert_eq!(graph.edges().len(), 1);
218 assert_eq!(graph.count_by_kind(EdgeKind::Defines), 1);
219 assert_eq!(graph.count_by_kind(EdgeKind::Calls), 0);
220 }
221
222 #[test]
223 fn find_callers_and_callees() {
224 let from_id = SymbolId::new(&camino::Utf8PathBuf::from("src/lib.rs"), "caller", 10);
225 let to_id = SymbolId::new(&camino::Utf8PathBuf::from("src/lib.rs"), "callee", 50);
226 let edges = vec![Edge {
227 from: from_id.clone(),
228 to: to_id.clone(),
229 kind: EdgeKind::Calls,
230 confidence: Confidence::Heuristic,
231 }];
232 let graph = Graph::new(edges);
233
234 let callers = graph.callers(&SymbolSelector::ById(to_id.clone()));
235 assert_eq!(callers.len(), 1);
236 assert_eq!(callers[0].from, from_id);
237
238 let callees = graph.callees(&SymbolSelector::ById(from_id));
239 assert_eq!(callees.len(), 1);
240 assert_eq!(callees[0].to, to_id);
241 }
242
243 #[test]
244 fn find_references_by_name() {
245 let from_id = SymbolId::new(&camino::Utf8PathBuf::from("src/lib.rs"), "caller", 10);
246 let to_id = SymbolId::new(&camino::Utf8PathBuf::from("src/lib.rs"), "callee", 50);
247 let edges = vec![Edge {
248 from: from_id.clone(),
249 to: to_id.clone(),
250 kind: EdgeKind::References,
251 confidence: Confidence::Heuristic,
252 }];
253 let graph = Graph::new(edges);
254
255 let refs = graph.find_references(&SymbolSelector::ById(to_id));
256 assert_eq!(refs.len(), 1);
257 }
258}