ryo_analysis/discovery/

engine.rs

//! DiscoveryEngine - Main symbol discovery interface.

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;

/// Main discovery engine for pattern-based symbol search.
///
/// # Example
/// ```rust,ignore
/// let engine = DiscoveryEngine::new(&code_graph, &registry, None, None);
///
/// let query = DiscoveryQuery::symbol("*Config")
///     .kinds(vec![SymbolKind::Struct])
///     .with_relations();
///
/// let result = engine.execute(&query)?;
/// for symbol in result.iter() {
///     println!("{}: {:?}", symbol.path, symbol.kind);
/// }
/// ```
pub struct DiscoveryEngine<'a> {
    graph: &'a CodeGraphV2,
    registry: &'a SymbolRegistry,
    _dataflow: Option<&'a DataFlowGraphV2>,
    typeflow: Option<&'a TypeFlowGraphV2>,
}

impl<'a> DiscoveryEngine<'a> {
    /// Create a new discovery engine.
    pub fn new(
        graph: &'a CodeGraphV2,
        registry: &'a SymbolRegistry,
        dataflow: Option<&'a DataFlowGraphV2>,
    ) -> Self {
        Self {
            graph,
            registry,
            _dataflow: dataflow,
            typeflow: None,
        }
    }

    /// Add TypeFlowGraphV2 to the engine.
    pub fn set_typeflow(mut self, typeflow: &'a TypeFlowGraphV2) -> Self {
        self.typeflow = Some(typeflow);
        self
    }

    /// Execute a discovery query.
    pub fn execute(&self, query: &DiscoveryQuery) -> DiscoveryResult {
        let mut result = DiscoveryResult::new();
        let mut matches = Vec::new();

        // Find matching symbols
        for (id, path) in self.registry.iter() {
            // Check pattern match
            if !query.pattern.matches(path.name()) {
                continue;
            }

            // Check crate filter
            if let Some(ref crate_name) = query.in_crate {
                if path.crate_name() != crate_name {
                    continue;
                }
            }

            // Check module filter
            if let Some(ref module) = query.in_module {
                if !path.to_string().contains(module) {
                    continue;
                }
            }

            // Check kind filter
            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;
                }
            }

            // Symbol matches!
            let kind = self.registry.kind(id).unwrap_or(SymbolKind::Other);
            let mut symbol = DiscoveredSymbol::new(id, path.clone(), kind);

            // Add persistent UUID if available
            if let Some(uuid) = self.registry.uuid(id) {
                symbol = symbol.with_uuid(uuid);
            }

            // Add metadata if available
            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());
            }

            // Calculate match score and reference count
            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);
        }

        // Apply type filter if specified and TypeFlowGraph is available
        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);
            }
        }

        // Sort based on query sort order
        match query.sort {
            SortOrder::Refs => {
                // Sort by reference count (descending), then by score
                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 => {
                // Alphabetical order
                matches.sort_by(|a, b| a.path.name().cmp(b.path.name()));
            }
            SortOrder::Impls => {
                // Sort by impl count (descending)
                matches.sort_by(|a, b| {
                    b.impl_count
                        .cmp(&a.impl_count)
                        .then_with(|| b.ref_count.cmp(&a.ref_count))
                });
            }
        }

        // Apply limit
        result.total_matches = matches.len();
        if let Some(limit) = query.limit {
            if matches.len() > limit {
                result.truncated = true;
                matches.truncate(limit);
            }
        }

        result.symbols = matches;

        // Build relation graph if requested
        if query.include_relations && !result.is_empty() {
            result.relations = Some(self.build_relations(&result, query.relation_depth));
        }

        result
    }

    /// Calculate match score for ranking.
    fn calculate_score(&self, name: &str, pattern: &crate::Pattern) -> f32 {
        let pattern_str = pattern.as_str();

        // Exact match gets highest score
        if name == pattern_str {
            return 1.0;
        }

        // Prefix match (pattern starts with name or vice versa)
        if pattern.is_exact() {
            return 0.0; // Should have matched exactly above
        }

        // For glob patterns, favor shorter names and prefix matches
        let base_score = 0.5;

        // Bonus for shorter names
        let length_bonus = 1.0 / (name.len() as f32 + 1.0) * 0.2;

        // Bonus for prefix match
        let prefix_bonus = if pattern_str.starts_with('*') {
            // Suffix pattern like "*Config" - favor exact suffix
            let suffix = pattern_str.trim_start_matches('*');
            if name.ends_with(suffix) {
                0.2
            } else {
                0.0
            }
        } else if pattern_str.ends_with('*') {
            // Prefix pattern like "get_*"
            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
    }

    /// Build relation graph for discovered symbols.
    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
    }

    /// Recursively add relations for a symbol.
    fn add_relations(
        &self,
        graph: &mut RelationGraph,
        id: crate::symbol::SymbolId,
        max_depth: usize,
        current_depth: usize,
    ) {
        if current_depth >= max_depth {
            return;
        }

        // Check if symbol exists in graph
        if !self.graph.contains(id) {
            return;
        }

        // Add outgoing relations
        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);
        }

        // Add incoming relations
        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);
        }

        // Add type reference relations from TypeFlow
        if let Some(typeflow) = self.typeflow {
            // Forward: what types does this symbol reference?
            for target_id in typeflow.types_used_by(id) {
                graph.add(id, target_id, RelationKind::TypeReferences);
            }
            // Reverse: who references this symbol as a type?
            for user_id in typeflow.type_users(id) {
                graph.add(id, user_id, RelationKind::TypeReferencedBy);
            }
        }
    }

    /// Find symbols by exact name.
    pub fn find_exact(&self, name: &str) -> DiscoveryResult {
        self.execute(&DiscoveryQuery::exact(name))
    }

    /// Find symbols matching a glob pattern.
    pub fn find_glob(&self, pattern: &str) -> DiscoveryResult {
        self.execute(&DiscoveryQuery::symbol(pattern))
    }

    /// Find all implementations of a trait.
    pub fn find_implementations(
        &self,
        trait_id: crate::symbol::SymbolId,
    ) -> Vec<crate::symbol::SymbolId> {
        self.graph.implementors_of(trait_id).collect()
    }

    /// Find all callers of a function.
    pub fn find_callers(&self, func_id: crate::symbol::SymbolId) -> Vec<crate::symbol::SymbolId> {
        self.graph.callers_of(func_id).collect()
    }

    /// Find all callees of a function.
    pub fn find_callees(&self, func_id: crate::symbol::SymbolId) -> Vec<crate::symbol::SymbolId> {
        self.graph.callees_of(func_id).collect()
    }

    // =========================================================================
    // TypeFlow filtering (V2 API)
    // =========================================================================

    /// Apply type filter to matched symbols using TypeFlowGraphV2.
    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()
    }

    /// Check if a symbol matches the type filter.
    fn matches_type_filter(
        &self,
        id: SymbolId,
        filter: &TypeFilter,
        typeflow: &TypeFlowGraphV2,
    ) -> bool {
        // Check return type filter
        if let Some(ref pattern) = filter.return_type {
            if !self.has_matching_usage_v2(id, UsageContext::ReturnType, pattern, typeflow) {
                return false;
            }
        }

        // Check parameter type filter
        if let Some(ref pattern) = filter.param_type {
            if !self.has_matching_usage_v2(id, UsageContext::ParamType, pattern, typeflow) {
                return false;
            }
        }

        // Check field type filter (structs/enums)
        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;
            }
        }

        // Check uses_type filter (any context)
        if let Some(ref pattern) = filter.uses_type {
            if !self.has_any_matching_usage_v2(id, pattern, typeflow) {
                return false;
            }
        }

        // Check has_bound filter (trait bounds)
        if let Some(ref pattern) = filter.has_bound {
            if !self.has_matching_bound_v2(id, pattern, typeflow) {
                return false;
            }
        }

        true
    }

    /// Check if there's a usage in a specific context matching the pattern (V2).
    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;
                }
                // V2 is String-free: resolve SymbolId to name via registry
                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
    }

    /// Check if there's any usage matching the pattern (any context) (V2).
    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) {
                // V2 is String-free: resolve SymbolId to name via registry
                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
    }

    /// Check if the symbol has trait bounds matching the pattern (V2).
    fn has_matching_bound_v2(
        &self,
        id: SymbolId,
        pattern: &Pattern,
        typeflow: &TypeFlowGraphV2,
    ) -> bool {
        // Get the definition node for this symbol
        let Some(def_node) = typeflow.definition(id) else {
            return false;
        };

        // Check trait bounds
        for bound_id in typeflow.trait_bounds(def_node) {
            if let Some(bound_data) = typeflow.get_trait_bound(bound_id) {
                // V2 is String-free: resolve SymbolId to name via registry
                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, &registry, None);

        let result = engine.find_glob("*Config");
        assert_eq!(result.len(), 3); // AppConfig, UserConfig, Config
    }

    #[test]
    fn test_find_with_crate_filter() {
        let (registry, graph) = setup();
        let engine = DiscoveryEngine::new(&graph, &registry, None);

        let query = DiscoveryQuery::symbol("*Config").in_crate("mylib");
        let result = engine.execute(&query);
        assert_eq!(result.len(), 2); // AppConfig, UserConfig
    }

    #[test]
    fn test_find_with_kind_filter() {
        let (registry, graph) = setup();
        let engine = DiscoveryEngine::new(&graph, &registry, None);

        let query = DiscoveryQuery::symbol("*").kind(SymbolKind::Function);
        let result = engine.execute(&query);
        assert_eq!(result.len(), 1); // handle
    }

    #[test]
    fn test_find_with_module_filter() {
        let (registry, graph) = setup();
        let engine = DiscoveryEngine::new(&graph, &registry, None);

        let query = DiscoveryQuery::symbol("*").in_module("handlers");
        let result = engine.execute(&query);
        assert_eq!(result.len(), 1); // handle
    }

    #[test]
    fn test_find_exact() {
        let (registry, graph) = setup();
        let engine = DiscoveryEngine::new(&graph, &registry, 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, &registry, 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);
    }
}