codesearch 0.1.12

//! Symbol Relationship Graph
//!
//! Tracks and queries relationships between symbols (inheritance, calls, references, etc.)

use super::{Symbol, SymbolRelation};
use super::SymbolRelationType;
use dashmap::DashMap;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

/// Relationship graph between symbols
#[derive(Clone)]
pub struct RelationshipGraph {
    /// Symbol ID to related symbols mapping
    relationships: Arc<DashMap<String, Vec<SymbolRelation>>>,

    /// Reverse relationships (for efficient lookup)
    reverse_relationships: Arc<DashMap<String, Vec<SymbolRelation>>>,

    /// All symbols (for context)
    symbols: Arc<DashMap<String, Symbol>>,
}

impl RelationshipGraph {
    /// Create a new relationship graph
    pub fn new() -> Self {
        Self {
            relationships: Arc::new(DashMap::new()),
            reverse_relationships: Arc::new(DashMap::new()),
            symbols: Arc::new(DashMap::new()),
        }
    }

    /// Add a symbol to the graph
    pub fn add_symbol(&self, symbol: Symbol) {
        self.symbols.insert(symbol.id.clone(), symbol);
    }

    /// Add a relationship between two symbols
    pub fn add_relationship(
        &self,
        from_symbol: &str,
        to_symbol: &str,
        relation_type: SymbolRelationType,
        confidence: f64,
    ) {
        let relation = SymbolRelation {
            symbol_id: to_symbol.to_string(),
            relation_type: relation_type.clone(),
            confidence,
        };

        // Add forward relationship
        self.relationships
            .entry(from_symbol.to_string())
            .or_insert_with(Vec::new)
            .push(relation.clone());

        // Add reverse relationship
        let reverse_type = relation_type.reverse();
        let reverse_relation = SymbolRelation {
            symbol_id: from_symbol.to_string(),
            relation_type: reverse_type,
            confidence,
        };

        self.reverse_relationships
            .entry(to_symbol.to_string())
            .or_insert_with(Vec::new)
            .push(reverse_relation);
    }

    /// Get all relationships for a symbol
    pub fn get_relationships(&self, symbol_id: &str) -> Vec<SymbolRelation> {
        self.relationships
            .get(symbol_id)
            .map(|relations| relations.value().clone())
            .unwrap_or_default()
    }

    /// Get relationships of a specific type
    pub fn get_relationships_by_type(
        &self,
        symbol_id: &str,
        relation_type: SymbolRelationType,
    ) -> Vec<SymbolRelation> {
        self.get_relationships(symbol_id)
            .into_iter()
            .filter(|r| r.relation_type == relation_type)
            .collect()
    }

    /// Get reverse relationships (what references this symbol)
    pub fn get_reverse_relationships(&self, symbol_id: &str) -> Vec<SymbolRelation> {
        self.reverse_relationships
            .get(symbol_id)
            .map(|relations| relations.value().clone())
            .unwrap_or_default()
    }

    /// Find symbols that inherit from the given symbol
    pub fn find_children(&self, symbol_id: &str) -> Vec<Symbol> {
        self.get_reverse_relationships(symbol_id)
            .into_iter()
            .filter(|r| r.relation_type == SymbolRelationType::Inherits)
            .filter_map(|r| self.symbols.get(&r.symbol_id).map(|s| s.value().clone()))
            .collect()
    }

    /// Find symbols that the given symbol inherits from
    pub fn find_parents(&self, symbol_id: &str) -> Vec<Symbol> {
        self.get_relationships(symbol_id)
            .into_iter()
            .filter(|r| r.relation_type == SymbolRelationType::Inherits)
            .filter_map(|r| self.symbols.get(&r.symbol_id).map(|s| s.value().clone()))
            .collect()
    }

    /// Find symbols called by the given symbol
    pub fn find_callees(&self, symbol_id: &str) -> Vec<Symbol> {
        self.get_relationships(symbol_id)
            .into_iter()
            .filter(|r| r.relation_type == SymbolRelationType::Calls)
            .filter_map(|r| self.symbols.get(&r.symbol_id).map(|s| s.value().clone()))
            .collect()
    }

    /// Find symbols that call the given symbol
    pub fn find_callers(&self, symbol_id: &str) -> Vec<Symbol> {
        self.get_reverse_relationships(symbol_id)
            .into_iter()
            .filter(|r| r.relation_type == SymbolRelationType::CalledBy)
            .filter_map(|r| self.symbols.get(&r.symbol_id).map(|s| s.value().clone()))
            .collect()
    }

    /// Find all symbols in the inheritance hierarchy
    pub fn find_hierarchy(&self, symbol_id: &str) -> Vec<Symbol> {
        let mut hierarchy = Vec::new();
        let mut visited = HashSet::new();

        self.collect_hierarchy(symbol_id, &mut hierarchy, &mut visited);

        hierarchy
    }

    fn collect_hierarchy(
        &self,
        symbol_id: &str,
        hierarchy: &mut Vec<Symbol>,
        visited: &mut HashSet<String>,
    ) {
        if visited.contains(symbol_id) {
            return;
        }

        visited.insert(symbol_id.to_string());

        if let Some(symbol) = self.symbols.get(symbol_id) {
            hierarchy.push(symbol.value().clone());

            // Add parents
            for parent in self.find_parents(symbol_id) {
                self.collect_hierarchy(&parent.id, hierarchy, visited);
            }

            // Add children
            for child in self.find_children(symbol_id) {
                self.collect_hierarchy(&child.id, hierarchy, visited);
            }
        }
    }

    /// Find related symbols by multiple relationship types
    pub fn find_related(
        &self,
        symbol_id: &str,
        relation_types: &[SymbolRelationType],
    ) -> Vec<Symbol> {
        let mut related = Vec::new();

        for relation in self.get_relationships(symbol_id) {
            if relation_types.contains(&relation.relation_type) {
                if let Some(symbol) = self.symbols.get(&relation.symbol_id) {
                    related.push(symbol.value().clone());
                }
            }
        }

        related
    }

    /// Get symbol by ID
    pub fn get_symbol(&self, id: &str) -> Option<Symbol> {
        self.symbols.get(id).map(|s| s.value().clone())
    }

    /// Get all symbols
    pub fn all_symbols(&self) -> Vec<Symbol> {
        self.symbols.iter().map(|e| e.value().clone()).collect()
    }

    /// Clear the graph
    pub fn clear(&self) {
        self.relationships.clear();
        self.reverse_relationships.clear();
        self.symbols.clear();
    }

    /// Build relationships from symbol analysis
    pub fn build_from_symbols(&self, symbols: &[Symbol]) {
        // Add all symbols to the graph
        for symbol in symbols {
            self.add_symbol(symbol.clone());
        }

        // Analyze relationships
        for symbol in symbols {
            self.analyze_symbol_relationships(symbol);
        }
    }

    /// Analyze relationships for a single symbol
    fn analyze_symbol_relationships(&self, symbol: &Symbol) {
        // Inheritance relationships
        if let Some(ref parent) = symbol.parent {
            self.add_relationship(
                &symbol.id,
                parent,
                SymbolRelationType::ContainedIn,
                1.0,
            );
        }

        // Type-based relationships
        if symbol.kind == super::SymbolKind::Class || symbol.kind == super::SymbolKind::Struct {
            // Look for inheritance in signature
            if symbol.signature.contains("extends") || symbol.signature.contains(":") {
                // Parse parent class from signature
                // This is simplified - real implementation would use proper parsing
            }
        }

        // Function/method relationships
        if symbol.kind == super::SymbolKind::Function || symbol.kind == super::SymbolKind::Method {
            self.analyze_function_calls(symbol);
        }
    }

    /// Analyze function calls within a symbol
    fn analyze_function_calls(&self, _symbol: &Symbol) {
        // In a real implementation, you'd:
        // 1. Read the function's code
        // 2. Parse function calls
        // 3. Add "Calls" relationships
        // For now, this is a placeholder
    }
}

impl Default for RelationshipGraph {
    fn default() -> Self {
        Self::new()
    }
}

impl SymbolRelationType {
    /// Get the reverse relationship type
    pub fn reverse(&self) -> SymbolRelationType {
        match self {
            SymbolRelationType::Inherits => SymbolRelationType::Inherits,
            SymbolRelationType::CalledBy => SymbolRelationType::Calls,
            SymbolRelationType::Implements => SymbolRelationType::CalledBy, // Simplified reverse
            SymbolRelationType::Calls => SymbolRelationType::CalledBy,
            SymbolRelationType::Contains => SymbolRelationType::ContainedIn,
            SymbolRelationType::ContainedIn => SymbolRelationType::Contains,
            SymbolRelationType::References => SymbolRelationType::ReferencedBy,
            SymbolRelationType::ReferencedBy => SymbolRelationType::References,
            SymbolRelationType::Overrides => SymbolRelationType::OverriddenBy,
            SymbolRelationType::OverriddenBy => SymbolRelationType::Overrides,
            SymbolRelationType::Instantiates => SymbolRelationType::InstantiatedBy,
            SymbolRelationType::InstantiatedBy => SymbolRelationType::Instantiates,
            SymbolRelationType::ParameterOf => SymbolRelationType::FieldOf, // Simplified reverse
            SymbolRelationType::ReturnTypeOf => SymbolRelationType::MethodOf, // Simplified reverse
            SymbolRelationType::FieldOf => SymbolRelationType::Contains, // Simplified reverse
            SymbolRelationType::MethodOf => SymbolRelationType::Contains, // Simplified reverse
            SymbolRelationType::Imports => SymbolRelationType::ImportedBy,
            SymbolRelationType::ImportedBy => SymbolRelationType::Imports,
            SymbolRelationType::Exports => SymbolRelationType::ExportedBy,
            SymbolRelationType::ExportedBy => SymbolRelationType::Exports,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    #[test]
    fn test_relationship_graph() {
        let graph = RelationshipGraph::new();

        // Add some symbols
        let parent = Symbol {
            id: "parent".to_string(),
            name: "ParentClass".to_string(),
            kind: super::super::SymbolKind::Class,
            file_path: "test.rs".to_string(),
            line: 10,
            column: 0,
            end_line: 15,
            signature: "".to_string(),
            documentation: None,
            visibility: super::super::SymbolVisibility::Public,
            parent: None,
            type_info: None,
            generics: vec![],
            annotations: vec![],
            attributes: vec![],
            metadata: HashMap::new(),
        };

        let child = Symbol {
            id: "child".to_string(),
            name: "ChildClass".to_string(),
            kind: super::super::SymbolKind::Class,
            file_path: "test.rs".to_string(),
            line: 20,
            column: 0,
            end_line: 25,
            signature: "".to_string(),
            documentation: None,
            visibility: super::super::SymbolVisibility::Public,
            parent: None,
            type_info: None,
            generics: vec![],
            annotations: vec![],
            attributes: vec![],
            metadata: HashMap::new(),
        };

        graph.add_symbol(parent.clone());
        graph.add_symbol(child.clone());

        // Add relationship
        graph.add_relationship(&child.id, &parent.id, SymbolRelationType::Inherits, 1.0);

        // Test relationships
        let parents = graph.find_parents(&child.id);
        assert_eq!(parents.len(), 1);
        assert_eq!(parents[0].name, "ParentClass");

        let children = graph.find_children(&parent.id);
        assert_eq!(children.len(), 1);
        assert_eq!(children[0].name, "ChildClass");
    }

    #[test]
    fn test_call_relationships() {
        let graph = RelationshipGraph::new();

        let caller = Symbol {
            id: "caller".to_string(),
            name: "caller_func".to_string(),
            kind: super::super::SymbolKind::Function,
            file_path: "test.rs".to_string(),
            line: 10,
            column: 0,
            end_line: 15,
            signature: "".to_string(),
            documentation: None,
            visibility: super::super::SymbolVisibility::Public,
            parent: None,
            type_info: None,
            generics: vec![],
            annotations: vec![],
            attributes: vec![],
            metadata: HashMap::new(),
        };

        let callee = Symbol {
            id: "callee".to_string(),
            name: "callee_func".to_string(),
            kind: super::super::SymbolKind::Function,
            file_path: "test.rs".to_string(),
            line: 20,
            column: 0,
            end_line: 25,
            signature: "".to_string(),
            documentation: None,
            visibility: super::super::SymbolVisibility::Public,
            parent: None,
            type_info: None,
            generics: vec![],
            annotations: vec![],
            attributes: vec![],
            metadata: HashMap::new(),
        };

        graph.add_symbol(caller.clone());
        graph.add_symbol(callee.clone());

        graph.add_relationship(&caller.id, &callee.id, SymbolRelationType::Calls, 1.0);

        // Test call relationships
        let callees = graph.find_callees(&caller.id);
        assert_eq!(callees.len(), 1);
        assert_eq!(callees[0].name, "callee_func");

        let callers = graph.find_callers(&callee.id);
        assert_eq!(callers.len(), 1);
        assert_eq!(callers[0].name, "caller_func");
    }
}