1use super::{CodeGraphV2, TypeFlowGraphV2};
14use crate::symbol::{SymbolId, SymbolPath, SymbolRegistry};
15use crate::Pattern;
16use crate::SymbolKind;
17
18pub struct QueryBuilder<'a> {
20 graph: &'a CodeGraphV2,
21 typeflow: &'a TypeFlowGraphV2,
22 registry: &'a SymbolRegistry,
23 filters: Vec<Filter>,
24}
25
26enum Filter {
28 Kind(SymbolKind),
29 KindOneOf(Vec<SymbolKind>),
30 Public,
31 CrateVisible,
32 Private,
33 InModule(String),
34 InCrate(String),
35 NamePattern(Pattern),
36 HasCallers,
37 HasNoCallers,
38 Implements(SymbolId),
39 UsedBy(SymbolId),
40}
41
42impl<'a> QueryBuilder<'a> {
43 pub fn new(
45 graph: &'a CodeGraphV2,
46 typeflow: &'a TypeFlowGraphV2,
47 registry: &'a SymbolRegistry,
48 ) -> Self {
49 Self {
50 graph,
51 typeflow,
52 registry,
53 filters: Vec::new(),
54 }
55 }
56
57 pub fn functions(mut self) -> Self {
61 self.filters.push(Filter::Kind(SymbolKind::Function));
62 self
63 }
64
65 pub fn structs(mut self) -> Self {
67 self.filters.push(Filter::Kind(SymbolKind::Struct));
68 self
69 }
70
71 pub fn enums(mut self) -> Self {
73 self.filters.push(Filter::Kind(SymbolKind::Enum));
74 self
75 }
76
77 pub fn traits(mut self) -> Self {
79 self.filters.push(Filter::Kind(SymbolKind::Trait));
80 self
81 }
82
83 pub fn modules(mut self) -> Self {
85 self.filters.push(Filter::Kind(SymbolKind::Mod));
86 self
87 }
88
89 pub fn impls(mut self) -> Self {
91 self.filters.push(Filter::Kind(SymbolKind::Impl));
92 self
93 }
94
95 pub fn methods(mut self) -> Self {
97 self.filters.push(Filter::Kind(SymbolKind::Method));
98 self
99 }
100
101 pub fn kind(mut self, kind: SymbolKind) -> Self {
103 self.filters.push(Filter::Kind(kind));
104 self
105 }
106
107 pub fn kinds(mut self, kinds: Vec<SymbolKind>) -> Self {
109 self.filters.push(Filter::KindOneOf(kinds));
110 self
111 }
112
113 pub fn public(mut self) -> Self {
117 self.filters.push(Filter::Public);
118 self
119 }
120
121 pub fn crate_visible(mut self) -> Self {
123 self.filters.push(Filter::CrateVisible);
124 self
125 }
126
127 pub fn private(mut self) -> Self {
129 self.filters.push(Filter::Private);
130 self
131 }
132
133 pub fn in_module(mut self, module: impl Into<String>) -> Self {
137 self.filters.push(Filter::InModule(module.into()));
138 self
139 }
140
141 pub fn in_crate(mut self, crate_name: impl Into<String>) -> Self {
143 self.filters.push(Filter::InCrate(crate_name.into()));
144 self
145 }
146
147 pub fn name_matches(mut self, pattern: Pattern) -> Self {
151 self.filters.push(Filter::NamePattern(pattern));
152 self
153 }
154
155 pub fn name(self, name: impl Into<String>) -> Self {
157 self.name_matches(Pattern::exact(name.into()))
158 }
159
160 pub fn name_glob(self, pattern: impl Into<String>) -> Self {
162 self.name_matches(Pattern::glob(pattern.into()))
163 }
164
165 pub fn has_callers(mut self) -> Self {
169 self.filters.push(Filter::HasCallers);
170 self
171 }
172
173 pub fn has_no_callers(mut self) -> Self {
175 self.filters.push(Filter::HasNoCallers);
176 self
177 }
178
179 pub fn implements_trait(mut self, trait_id: SymbolId) -> Self {
181 self.filters.push(Filter::Implements(trait_id));
182 self
183 }
184
185 pub fn used_by(mut self, user_id: SymbolId) -> Self {
187 self.filters.push(Filter::UsedBy(user_id));
188 self
189 }
190
191 pub fn collect(self) -> Vec<SymbolId> {
195 self.registry
196 .iter()
197 .map(|(id, _)| id)
198 .filter(|&id| self.matches(id))
199 .collect()
200 }
201
202 pub fn first(self) -> Option<SymbolId> {
204 self.registry
205 .iter()
206 .map(|(id, _)| id)
207 .find(|&id| self.matches(id))
208 }
209
210 pub fn count(self) -> usize {
212 self.registry
213 .iter()
214 .map(|(id, _)| id)
215 .filter(|&id| self.matches(id))
216 .count()
217 }
218
219 pub fn exists(self) -> bool {
221 self.registry
222 .iter()
223 .map(|(id, _)| id)
224 .any(|id| self.matches(id))
225 }
226
227 fn matches(&self, id: SymbolId) -> bool {
229 let path = match self.registry.resolve(id) {
230 Some(p) => p,
231 None => return false,
232 };
233
234 for filter in &self.filters {
235 if !self.matches_filter(id, path, filter) {
236 return false;
237 }
238 }
239
240 true
241 }
242
243 fn matches_filter(&self, id: SymbolId, path: &SymbolPath, filter: &Filter) -> bool {
245 match filter {
246 Filter::Kind(kind) => self
247 .registry
248 .kind(id)
249 .is_some_and(|k| k == *kind || k.matches(kind)),
250 Filter::KindOneOf(kinds) => self
251 .registry
252 .kind(id)
253 .is_some_and(|k| kinds.iter().any(|kind| k.matches(kind))),
254 Filter::Public => self.registry.visibility(id).is_some_and(|v| v.is_public()),
255 Filter::CrateVisible => self
256 .registry
257 .visibility(id)
258 .is_some_and(|v| v.is_crate_visible()),
259 Filter::Private => self.registry.visibility(id).is_none_or(|v| v.is_private()),
260 Filter::InModule(module) => path.to_string().contains(module),
261 Filter::InCrate(crate_name) => path.crate_name() == crate_name,
262 Filter::NamePattern(pattern) => pattern.matches(path.name()),
263 Filter::HasCallers => self.graph.callers_of(id).next().is_some(),
264 Filter::HasNoCallers => self.graph.callers_of(id).next().is_none(),
265 Filter::Implements(trait_id) => self.graph.implementors_of(*trait_id).any(|x| x == id),
266 Filter::UsedBy(user_id) => self.typeflow.types_used_by(*user_id).any(|x| x == id),
267 }
268 }
269}
270
271#[cfg(test)]
272mod tests {
273 use super::*;
274 use crate::symbol::Visibility;
275
276 trait QueryExt {
277 fn query<'a>(
278 &'a self,
279 typeflow: &'a TypeFlowGraphV2,
280 registry: &'a SymbolRegistry,
281 ) -> QueryBuilder<'a>;
282 }
283
284 impl QueryExt for CodeGraphV2 {
285 fn query<'a>(
286 &'a self,
287 typeflow: &'a TypeFlowGraphV2,
288 registry: &'a SymbolRegistry,
289 ) -> QueryBuilder<'a> {
290 QueryBuilder::new(self, typeflow, registry)
291 }
292 }
293
294 fn setup() -> (SymbolRegistry, CodeGraphV2, TypeFlowGraphV2) {
295 let mut registry = SymbolRegistry::new();
296
297 let func1 = registry
298 .register_with_metadata(
299 SymbolPath::parse("mylib::handlers::handle").unwrap(),
300 SymbolKind::Function,
301 None,
302 Some(Visibility::Public),
303 )
304 .unwrap();
305 let func2 = registry
306 .register_with_metadata(
307 SymbolPath::parse("mylib::handlers::process").unwrap(),
308 SymbolKind::Function,
309 None,
310 Some(Visibility::Private),
311 )
312 .unwrap();
313 let struct1 = registry
314 .register_with_metadata(
315 SymbolPath::parse("mylib::models::User").unwrap(),
316 SymbolKind::Struct,
317 None,
318 Some(Visibility::Public),
319 )
320 .unwrap();
321
322 let mut graph = CodeGraphV2::new();
323 graph.add_node(func1);
324 graph.add_node(func2);
325 graph.add_node(struct1);
326 graph.add_to_kind_index(func1, SymbolKind::Function);
327 graph.add_to_kind_index(func2, SymbolKind::Function);
328 graph.add_to_kind_index(struct1, SymbolKind::Struct);
329 graph.add_edge(func1, func2, super::super::CodeEdgeV2::Calls);
330
331 let typeflow = TypeFlowGraphV2::new();
332 (registry, graph, typeflow)
333 }
334
335 #[test]
336 fn test_query_functions() {
337 let (registry, graph, typeflow) = setup();
338 let results = graph.query(&typeflow, ®istry).functions().collect();
339 assert_eq!(results.len(), 2);
340 }
341
342 #[test]
343 fn test_query_structs() {
344 let (registry, graph, typeflow) = setup();
345 let results = graph.query(&typeflow, ®istry).structs().collect();
346 assert_eq!(results.len(), 1);
347 }
348
349 #[test]
350 fn test_query_public() {
351 let (registry, graph, typeflow) = setup();
352 let results = graph.query(&typeflow, ®istry).public().collect();
353 assert_eq!(results.len(), 2); }
355
356 #[test]
357 fn test_query_in_module() {
358 let (registry, graph, typeflow) = setup();
359 let results = graph
360 .query(&typeflow, ®istry)
361 .in_module("handlers")
362 .collect();
363 assert_eq!(results.len(), 2);
364 }
365
366 #[test]
367 fn test_query_combined() {
368 let (registry, graph, typeflow) = setup();
369 let results = graph
370 .query(&typeflow, ®istry)
371 .functions()
372 .public()
373 .in_module("handlers")
374 .collect();
375 assert_eq!(results.len(), 1);
376 }
377
378 #[test]
379 fn test_query_name_glob() {
380 let (registry, graph, typeflow) = setup();
381 let results = graph.query(&typeflow, ®istry).name_glob("*er").collect();
382 assert_eq!(results.len(), 1);
384 }
385
386 #[test]
387 fn test_query_has_callers() {
388 let (registry, graph, typeflow) = setup();
389 let results = graph
390 .query(&typeflow, ®istry)
391 .functions()
392 .has_callers()
393 .collect();
394 assert_eq!(results.len(), 1);
396 }
397
398 #[test]
399 fn test_query_count() {
400 let (registry, graph, typeflow) = setup();
401 let count = graph.query(&typeflow, ®istry).functions().count();
402 assert_eq!(count, 2);
403 }
404
405 #[test]
406 fn test_query_exists() {
407 let (registry, graph, typeflow) = setup();
408 assert!(graph.query(&typeflow, ®istry).structs().exists());
409 assert!(!graph.query(&typeflow, ®istry).traits().exists());
410 }
411}