ryo_analysis/query/

derive_index.rs

//! DeriveIndex: Fast lookup for derive trait validation.
//!
//! This index pre-computes the relationship between types and their derive traits,
//! enabling O(1) lookup for derive possibility checks instead of O(N) AST traversal.

use super::{CodeGraphV2, TypeFlowGraphV2};
use crate::ast::ASTRegistry;
use crate::symbol::SymbolRegistry;
use crate::SymbolId;
use ryo_source::pure::{PureAttrMeta, PureItem};
use serde::Serialize;
use slotmap::SecondaryMap;
use smallvec::SmallVec;

/// Index for fast derive trait validation.
///
/// # Memory Layout
/// ```text
/// DeriveIndex
/// ├── symbol_derives: SecondaryMap<SymbolId, SmallVec<[String; 4]>>
/// │   └── struct/enum → derive trait names (e.g., ["Debug", "Clone"])
/// └── field_type_names: SecondaryMap<SymbolId, SmallVec<[String; 8]>>
///     └── struct/enum → field type names (for derive validation)
/// ```
#[derive(Clone, Default, Debug, Serialize)]
pub struct DeriveIndex {
    /// struct/enum SymbolId → derive trait names.
    /// Most types derive 2-4 traits, so SmallVec<[String; 4]> is optimal.
    symbol_derives: SecondaryMap<SymbolId, SmallVec<[String; 4]>>,

    /// struct/enum SymbolId → field type names.
    /// Used for checking if all fields implement the trait.
    field_type_names: SecondaryMap<SymbolId, SmallVec<[String; 8]>>,
}

impl DeriveIndex {
    /// Create a new empty index.
    pub fn new() -> Self {
        Self::default()
    }

    /// Build the index from ASTRegistry, CodeGraph, TypeFlow, and SymbolRegistry.
    pub fn build(
        ast_registry: &ASTRegistry,
        code_graph: &CodeGraphV2,
        typeflow: &TypeFlowGraphV2,
        symbol_registry: &SymbolRegistry,
    ) -> Self {
        let mut index = Self::new();
        index.rebuild_all(ast_registry, code_graph, typeflow, symbol_registry);
        index
    }

    /// Rebuild the entire index.
    pub fn rebuild_all(
        &mut self,
        ast_registry: &ASTRegistry,
        code_graph: &CodeGraphV2,
        typeflow: &TypeFlowGraphV2,
        symbol_registry: &SymbolRegistry,
    ) {
        self.symbol_derives.clear();
        self.field_type_names.clear();

        for (id, item) in ast_registry.iter() {
            self.index_item(id, item, code_graph, typeflow, symbol_registry);
        }
    }

    /// Incrementally update for specific symbols.
    ///
    /// This is O(S) where S is the number of affected symbols,
    /// instead of O(N) for full rebuild.
    pub fn rebuild_for_symbols(
        &mut self,
        symbols: &[SymbolId],
        ast_registry: &ASTRegistry,
        code_graph: &CodeGraphV2,
        typeflow: &TypeFlowGraphV2,
        symbol_registry: &SymbolRegistry,
    ) {
        for &symbol_id in symbols {
            // Remove old entries
            self.symbol_derives.remove(symbol_id);
            self.field_type_names.remove(symbol_id);

            // Re-index if the symbol still exists
            if let Some(item) = ast_registry.get(symbol_id) {
                self.index_item(symbol_id, item, code_graph, typeflow, symbol_registry);
            }
        }
    }

    /// Index a single item.
    fn index_item(
        &mut self,
        id: SymbolId,
        item: &PureItem,
        code_graph: &CodeGraphV2,
        typeflow: &TypeFlowGraphV2,
        symbol_registry: &SymbolRegistry,
    ) {
        let attrs = match item {
            PureItem::Struct(s) => &s.attrs,
            PureItem::Enum(e) => &e.attrs,
            _ => return,
        };

        // Extract derive traits
        let mut derives: SmallVec<[String; 4]> = SmallVec::new();
        for attr in attrs {
            if attr.path == "derive" {
                if let PureAttrMeta::List(args) = &attr.meta {
                    for trait_name in args.split(',').map(|s| s.trim()) {
                        if !trait_name.is_empty() {
                            derives.push(trait_name.to_string());
                        }
                    }
                }
            }
        }

        if !derives.is_empty() {
            self.symbol_derives.insert(id, derives);
        }

        // Extract field type names from TypeFlow
        let mut field_types: SmallVec<[String; 8]> = SmallVec::new();
        for child_id in code_graph.children_of(id) {
            // Get the type that this field uses (via TypeFlow)
            for use_id in typeflow.types_used_by(child_id) {
                if let Some(path) = symbol_registry.resolve(use_id) {
                    field_types.push(path.name().to_string());
                    break; // Only first type reference
                }
            }
        }

        if !field_types.is_empty() {
            self.field_type_names.insert(id, field_types);
        }
    }

    // ========================================================================
    // Query Methods
    // ========================================================================

    /// Iterate over all symbols and their derives.
    pub fn iter_derives(&self) -> impl Iterator<Item = (SymbolId, &SmallVec<[String; 4]>)> {
        self.symbol_derives.iter()
    }

    /// Get derive traits for a symbol.
    pub fn get_derives(&self, id: SymbolId) -> Option<&SmallVec<[String; 4]>> {
        self.symbol_derives.get(id)
    }

    /// Get field type names for a symbol.
    pub fn get_field_types(&self, id: SymbolId) -> Option<&SmallVec<[String; 8]>> {
        self.field_type_names.get(id)
    }

    /// Check if a symbol has a specific derive trait.
    pub fn has_derive(&self, id: SymbolId, trait_name: &str) -> bool {
        self.symbol_derives
            .get(id)
            .map(|derives| derives.iter().any(|d| d == trait_name))
            .unwrap_or(false)
    }

    /// Get all symbols that derive a specific trait.
    pub fn symbols_deriving(&self, trait_name: &str) -> Vec<SymbolId> {
        self.symbol_derives
            .iter()
            .filter(|(_, derives)| derives.iter().any(|d| d == trait_name))
            .map(|(id, _)| id)
            .collect()
    }

    /// Get statistics about the index.
    pub fn stats(&self) -> DeriveIndexStats {
        let total_derives: usize = self.symbol_derives.values().map(|v| v.len()).sum();
        let total_fields: usize = self.field_type_names.values().map(|v| v.len()).sum();

        DeriveIndexStats {
            symbols_with_derives: self.symbol_derives.len(),
            total_derives,
            symbols_with_fields: self.field_type_names.len(),
            total_field_types: total_fields,
        }
    }
}

/// Statistics about the DeriveIndex.
#[derive(Debug, Clone)]
pub struct DeriveIndexStats {
    /// Number of distinct symbols that carry at least one `#[derive(...)]`.
    pub symbols_with_derives: usize,
    /// Sum of derive entries across all symbols (counts duplicates per
    /// symbol).
    pub total_derives: usize,
    /// Number of distinct symbols that have at least one indexed field
    /// type.
    pub symbols_with_fields: usize,
    /// Sum of indexed field-type entries across all symbols.
    pub total_field_types: usize,
}

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

    #[test]
    fn test_derive_index_creation() {
        let index = DeriveIndex::new();
        assert_eq!(index.stats().symbols_with_derives, 0);
    }
}