use super::{
DiscoveredSymbol, DiscoveryQuery, DiscoveryResult, RelationGraph, RelationKind, SortOrder,
TypeFilter,
};
use crate::query::{CodeEdgeV2, CodeGraphV2, DataFlowGraphV2, TypeFlowGraphV2, UsageContext};
use crate::symbol::{SymbolId, SymbolRegistry};
use crate::Pattern;
use crate::SymbolKind;
pub struct DiscoveryEngine<'a> {
graph: &'a CodeGraphV2,
registry: &'a SymbolRegistry,
_dataflow: Option<&'a DataFlowGraphV2>,
typeflow: Option<&'a TypeFlowGraphV2>,
}
impl<'a> DiscoveryEngine<'a> {
pub fn new(
graph: &'a CodeGraphV2,
registry: &'a SymbolRegistry,
dataflow: Option<&'a DataFlowGraphV2>,
) -> Self {
Self {
graph,
registry,
_dataflow: dataflow,
typeflow: None,
}
}
pub fn set_typeflow(mut self, typeflow: &'a TypeFlowGraphV2) -> Self {
self.typeflow = Some(typeflow);
self
}
pub fn execute(&self, query: &DiscoveryQuery) -> DiscoveryResult {
let mut result = DiscoveryResult::new();
let mut matches = Vec::new();
for (id, path) in self.registry.iter() {
if !query.pattern.matches(path.name()) {
continue;
}
if let Some(ref crate_name) = query.in_crate {
if path.crate_name() != crate_name {
continue;
}
}
if let Some(ref module) = query.in_module {
if !path.to_string().contains(module) {
continue;
}
}
if let Some(ref kinds) = query.kinds {
if let Some(kind) = self.registry.kind(id) {
if !kinds.iter().any(|k| k.matches(&kind)) {
continue;
}
} else {
continue;
}
}
let kind = self.registry.kind(id).unwrap_or(SymbolKind::Other);
let mut symbol = DiscoveredSymbol::new(id, path.clone(), kind);
if let Some(uuid) = self.registry.uuid(id) {
symbol = symbol.with_uuid(uuid);
}
if let Some(span) = self.registry.span(id) {
symbol = symbol.with_span(span.clone());
}
if let Some(vis) = self.registry.visibility(id) {
symbol = symbol.with_visibility(vis.clone());
}
let score = self.calculate_score(path.name(), &query.pattern);
let ref_count = self.graph.reference_count(id);
let impl_count = self.graph.impl_count(id);
symbol = symbol
.with_score(score)
.with_ref_count(ref_count)
.with_impl_count(impl_count);
matches.push(symbol);
}
if let (Some(ref type_filter), Some(typeflow)) = (&query.type_filter, self.typeflow) {
if !type_filter.is_empty() {
matches = self.apply_type_filter(matches, type_filter, typeflow);
}
}
match query.sort {
SortOrder::Refs => {
matches.sort_by(|a, b| {
b.ref_count.cmp(&a.ref_count).then_with(|| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
})
});
}
SortOrder::Alpha => {
matches.sort_by(|a, b| a.path.name().cmp(b.path.name()));
}
SortOrder::Impls => {
matches.sort_by(|a, b| {
b.impl_count
.cmp(&a.impl_count)
.then_with(|| b.ref_count.cmp(&a.ref_count))
});
}
}
result.total_matches = matches.len();
if let Some(limit) = query.limit {
if matches.len() > limit {
result.truncated = true;
matches.truncate(limit);
}
}
result.symbols = matches;
if query.include_relations && !result.is_empty() {
result.relations = Some(self.build_relations(&result, query.relation_depth));
}
result
}
fn calculate_score(&self, name: &str, pattern: &crate::Pattern) -> f32 {
let pattern_str = pattern.as_str();
if name == pattern_str {
return 1.0;
}
if pattern.is_exact() {
return 0.0; }
let base_score = 0.5;
let length_bonus = 1.0 / (name.len() as f32 + 1.0) * 0.2;
let prefix_bonus = if pattern_str.starts_with('*') {
let suffix = pattern_str.trim_start_matches('*');
if name.ends_with(suffix) {
0.2
} else {
0.0
}
} else if pattern_str.ends_with('*') {
let prefix = pattern_str.trim_end_matches('*');
if name.starts_with(prefix) {
0.2
} else {
0.0
}
} else {
0.0
};
base_score + length_bonus + prefix_bonus
}
fn build_relations(&self, result: &DiscoveryResult, max_depth: usize) -> RelationGraph {
let mut graph = RelationGraph::new();
for symbol in result.iter() {
self.add_relations(&mut graph, symbol.id, max_depth, 0);
}
graph
}
fn add_relations(
&self,
graph: &mut RelationGraph,
id: crate::symbol::SymbolId,
max_depth: usize,
current_depth: usize,
) {
if current_depth >= max_depth {
return;
}
if !self.graph.contains(id) {
return;
}
for edge_data in self.graph.outgoing_edges(id) {
let target_id = edge_data.to;
let kind = match edge_data.kind {
CodeEdgeV2::Calls => RelationKind::Calls,
CodeEdgeV2::Implements => RelationKind::Implements,
CodeEdgeV2::Contains => RelationKind::Contains,
};
graph.add(id, target_id, kind);
}
for edge_data in self.graph.incoming_edges(id) {
let source_id = edge_data.from;
let kind = match edge_data.kind {
CodeEdgeV2::Calls => RelationKind::CalledBy,
CodeEdgeV2::Implements => RelationKind::ImplementedBy,
CodeEdgeV2::Contains => RelationKind::ContainedBy,
};
graph.add(id, source_id, kind);
}
if let Some(typeflow) = self.typeflow {
for target_id in typeflow.types_used_by(id) {
graph.add(id, target_id, RelationKind::TypeReferences);
}
for user_id in typeflow.type_users(id) {
graph.add(id, user_id, RelationKind::TypeReferencedBy);
}
}
}
pub fn find_exact(&self, name: &str) -> DiscoveryResult {
self.execute(&DiscoveryQuery::exact(name))
}
pub fn find_glob(&self, pattern: &str) -> DiscoveryResult {
self.execute(&DiscoveryQuery::symbol(pattern))
}
pub fn find_implementations(
&self,
trait_id: crate::symbol::SymbolId,
) -> Vec<crate::symbol::SymbolId> {
self.graph.implementors_of(trait_id).collect()
}
pub fn find_callers(&self, func_id: crate::symbol::SymbolId) -> Vec<crate::symbol::SymbolId> {
self.graph.callers_of(func_id).collect()
}
pub fn find_callees(&self, func_id: crate::symbol::SymbolId) -> Vec<crate::symbol::SymbolId> {
self.graph.callees_of(func_id).collect()
}
fn apply_type_filter(
&self,
symbols: Vec<DiscoveredSymbol>,
filter: &TypeFilter,
typeflow: &TypeFlowGraphV2,
) -> Vec<DiscoveredSymbol> {
symbols
.into_iter()
.filter(|symbol| self.matches_type_filter(symbol.id, filter, typeflow))
.collect()
}
fn matches_type_filter(
&self,
id: SymbolId,
filter: &TypeFilter,
typeflow: &TypeFlowGraphV2,
) -> bool {
if let Some(ref pattern) = filter.return_type {
if !self.has_matching_usage_v2(id, UsageContext::ReturnType, pattern, typeflow) {
return false;
}
}
if let Some(ref pattern) = filter.param_type {
if !self.has_matching_usage_v2(id, UsageContext::ParamType, pattern, typeflow) {
return false;
}
}
if let Some(ref pattern) = filter.field_type {
let has_field =
self.has_matching_usage_v2(id, UsageContext::FieldType, pattern, typeflow)
|| self.has_matching_usage_v2(
id,
UsageContext::VariantField,
pattern,
typeflow,
);
if !has_field {
return false;
}
}
if let Some(ref pattern) = filter.uses_type {
if !self.has_any_matching_usage_v2(id, pattern, typeflow) {
return false;
}
}
if let Some(ref pattern) = filter.has_bound {
if !self.has_matching_bound_v2(id, pattern, typeflow) {
return false;
}
}
true
}
fn has_matching_usage_v2(
&self,
id: SymbolId,
expected_context: UsageContext,
pattern: &Pattern,
typeflow: &TypeFlowGraphV2,
) -> bool {
for usage_id in typeflow.usages(id) {
if let Some(usage_data) = typeflow.get_usage(usage_id) {
if usage_data.context != expected_context {
continue;
}
if let Some(resolved_id) = usage_data.resolved {
if let Some(path) = self.registry.resolve(resolved_id) {
if pattern.matches(path.name()) {
return true;
}
}
}
}
}
false
}
fn has_any_matching_usage_v2(
&self,
id: SymbolId,
pattern: &Pattern,
typeflow: &TypeFlowGraphV2,
) -> bool {
for usage_id in typeflow.usages(id) {
if let Some(usage_data) = typeflow.get_usage(usage_id) {
if let Some(resolved_id) = usage_data.resolved {
if let Some(path) = self.registry.resolve(resolved_id) {
if pattern.matches(path.name()) {
return true;
}
}
}
}
}
false
}
fn has_matching_bound_v2(
&self,
id: SymbolId,
pattern: &Pattern,
typeflow: &TypeFlowGraphV2,
) -> bool {
let Some(def_node) = typeflow.definition(id) else {
return false;
};
for bound_id in typeflow.trait_bounds(def_node) {
if let Some(bound_data) = typeflow.get_trait_bound(bound_id) {
if let Some(trait_id) = bound_data.trait_id {
if let Some(path) = self.registry.resolve(trait_id) {
if pattern.matches(path.name()) {
return true;
}
}
}
}
}
false
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::symbol::SymbolPath;
use crate::SymbolKind;
fn setup() -> (SymbolRegistry, CodeGraphV2) {
let mut registry = SymbolRegistry::new();
registry
.register(
SymbolPath::parse("mylib::AppConfig").unwrap(),
SymbolKind::Struct,
)
.unwrap();
registry
.register(
SymbolPath::parse("mylib::UserConfig").unwrap(),
SymbolKind::Struct,
)
.unwrap();
registry
.register(
SymbolPath::parse("mylib::handlers::handle").unwrap(),
SymbolKind::Function,
)
.unwrap();
registry
.register(
SymbolPath::parse("other::Config").unwrap(),
SymbolKind::Struct,
)
.unwrap();
let graph = CodeGraphV2::new();
(registry, graph)
}
#[test]
fn test_find_by_pattern() {
let (registry, graph) = setup();
let engine = DiscoveryEngine::new(&graph, ®istry, None);
let result = engine.find_glob("*Config");
assert_eq!(result.len(), 3); }
#[test]
fn test_find_with_crate_filter() {
let (registry, graph) = setup();
let engine = DiscoveryEngine::new(&graph, ®istry, None);
let query = DiscoveryQuery::symbol("*Config").in_crate("mylib");
let result = engine.execute(&query);
assert_eq!(result.len(), 2); }
#[test]
fn test_find_with_kind_filter() {
let (registry, graph) = setup();
let engine = DiscoveryEngine::new(&graph, ®istry, None);
let query = DiscoveryQuery::symbol("*").kind(SymbolKind::Function);
let result = engine.execute(&query);
assert_eq!(result.len(), 1); }
#[test]
fn test_find_with_module_filter() {
let (registry, graph) = setup();
let engine = DiscoveryEngine::new(&graph, ®istry, None);
let query = DiscoveryQuery::symbol("*").in_module("handlers");
let result = engine.execute(&query);
assert_eq!(result.len(), 1); }
#[test]
fn test_find_exact() {
let (registry, graph) = setup();
let engine = DiscoveryEngine::new(&graph, ®istry, None);
let result = engine.find_exact("AppConfig");
assert_eq!(result.len(), 1);
assert_eq!(result.first().unwrap().path.name(), "AppConfig");
}
#[test]
fn test_find_with_limit() {
let (registry, graph) = setup();
let engine = DiscoveryEngine::new(&graph, ®istry, None);
let query = DiscoveryQuery::symbol("*").limit(2);
let result = engine.execute(&query);
assert_eq!(result.len(), 2);
assert!(result.truncated);
assert_eq!(result.total_matches, 4);
}
}