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, StdImplCache, TypeFlowGraphV2};
use crate::ast::ASTRegistry;
use crate::symbol::SymbolRegistry;
use crate::SymbolId;
use ryo_source::pure::{PureAttrMeta, PureFields, PureItem, PureType};
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.
        //
        // Bug fix (precheck-fix-1): `code_graph.children_of(struct_id)` returns
        // ALL graph children — for `pub struct StmtConverter;` this includes
        // every method defined in `impl StmtConverter { ... }`. Pulling their
        // return types in as "field types" caused `verify_precheck` to
        // fabricate phantom fields (e.g. `[Result, Result, Result, Result]`)
        // for unit structs and fail their `Default` derive check, which in
        // turn rejected every speculative mutation at baseline.
        //
        // Filter to children whose `SymbolKind` is actually `Field` so we
        // only walk struct fields (and enum variant fields, via SymbolKind),
        // not methods or other associated items.
        let mut field_types: SmallVec<[String; 8]> = SmallVec::new();
        for child_id in code_graph.children_of(id) {
            // Only consider true fields — skip methods, consts, associated
            // types, and any other non-field children.
            if !matches!(
                symbol_registry.kind(child_id),
                Some(crate::SymbolKind::Field)
            ) {
                continue;
            }
            // 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
                }
            }
        }

        // Augment with field type names extracted directly from the AST.
        //
        // The typeflow-driven path above never reports primitives
        // (`builder_typeflow_v2.rs:906-907` deliberately skips them) and rarely
        // reports std containers (`Vec`, `HashMap`, …) because their
        // `SymbolId` is not registered for arbitrary user crates.
        //
        // Additionally, the typeflow path may fail to resolve user-defined
        // types that are defined in the same file but whose `SymbolId` was not
        // registered in the typeflow graph at indexing time (e.g. a bare
        // `pub struct Bar;` used as a field type in the same crate).
        //
        // The AST-direct augmentation below closes both gaps:
        // - For primitive / std-container heads: push via `walk_type_for_std_heads`.
        // - For user-defined heads (not primitive, not std-container): push via
        //   `walk_type_for_user_heads` so that `verify_precheck` can apply the
        //   conservative user-defined-type check on them as well.
        // Deduplication is applied in both cases.
        let std_impls = StdImplCache::default();
        let mut push_if_relevant = |field_types: &mut SmallVec<[String; 8]>, ty: &PureType| {
            // Primitive / std-container heads.
            walk_type_for_std_heads(ty, &std_impls, &mut |head| {
                if !field_types.iter().any(|existing| existing == &head) {
                    field_types.push(head);
                }
            });
            // User-defined heads (complement to typeflow path).
            walk_type_for_user_heads(ty, &std_impls, &mut |head| {
                if !field_types.iter().any(|existing| existing == &head) {
                    field_types.push(head);
                }
            });
        };
        match item {
            PureItem::Struct(s) => {
                collect_field_types_from_fields(&s.fields, &mut field_types, &mut push_if_relevant)
            }
            PureItem::Enum(e) => {
                for variant in &e.variants {
                    collect_field_types_from_fields(
                        &variant.fields,
                        &mut field_types,
                        &mut push_if_relevant,
                    );
                }
            }
            _ => {}
        }

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

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

    /// Test-only helper: directly insert a `(symbol → derives, field_types)`
    /// entry, bypassing the typeflow-driven `rebuild_*` path.
    ///
    /// This exists so callers can construct synthetic inputs that exercise
    /// downstream consumers of the index (notably
    /// `ryo_verification::GraphVerifier::verify_precheck`) at boundary
    /// conditions that the natural build path filters out — for example a
    /// `#[derive(Hash)]` struct whose field type is a primitive (`f32`), which
    /// `builder_typeflow_v2.rs` skips during typeflow population
    /// (`// Skip primitives - only add usages for user-defined types`) and
    /// therefore never reaches the verifier's reject branch organically.
    ///
    /// Production code MUST NOT call this. The method is `#[doc(hidden)]` to
    /// keep it out of the rendered API surface.
    #[doc(hidden)]
    pub fn insert_for_test(
        &mut self,
        id: SymbolId,
        derives: impl IntoIterator<Item = String>,
        field_types: impl IntoIterator<Item = String>,
    ) {
        let derives: SmallVec<[String; 4]> = derives.into_iter().collect();
        let field_types: SmallVec<[String; 8]> = field_types.into_iter().collect();
        if !derives.is_empty() {
            self.symbol_derives.insert(id, derives);
        }
        if !field_types.is_empty() {
            self.field_type_names.insert(id, field_types);
        }
    }

    /// 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,
        }
    }
}

/// Walk a `PureFields` and feed each field's `PureType` to `push`.
///
/// Handles the three field shapes — named (`{ x: T }`), tuple (`(T, U)`), and
/// unit — so callers don't have to repeat the match. Used by `index_item` for
/// both struct fields and enum variant fields.
fn collect_field_types_from_fields(
    fields: &PureFields,
    field_types: &mut SmallVec<[String; 8]>,
    push: &mut dyn FnMut(&mut SmallVec<[String; 8]>, &PureType),
) {
    match fields {
        PureFields::Named(named) => {
            for f in named {
                push(field_types, &f.ty);
            }
        }
        PureFields::Tuple(tuple_fields) => {
            for f in tuple_fields {
                push(field_types, &f.ty);
            }
        }
        PureFields::Unit => {}
    }
}

/// Walk a `PureType`, surfacing each primitive- or std-container-typed sub-term
/// to `sink`. Recurses into the structural composite variants so a `&f32` or
/// `(f32, u32)` field has every relevant inner type reach the derive verifier.
///
/// Variants handled:
/// - `PureType::Path(p)` — extracts the head name (`p` stripped of generics
///   and module path), pushes when `StdImplCache::is_primitive` or
///   `is_std_container` accepts it. Examples:
///   - `f32` → `f32`
///   - `std::collections::HashMap<K, V>` → `HashMap` (note: the head walk does
///     **not** descend into the generic arguments encoded inside the string;
///     see [Limitations] below).
///   - `my_crate::Config` → nothing (user-defined; the typeflow path covers
///     these via `SymbolId`).
/// - `PureType::Ref { ty, .. }` — recurses into the referent (`&T: Hash`
///   requires `T: Hash`, so primitives behind a reference still matter).
/// - `PureType::Tuple(types)` — recurses into each element (`(A, B): Hash`
///   requires every element to impl `Hash`).
/// - `PureType::Array { ty, .. }` / `PureType::Slice(ty)` — recurses into the
///   element type.
/// - Other variants (`Fn`, `ImplTrait`, `TraitObject`, `Infer`, `Never`,
///   `Other`) are deliberately skipped: they either can't appear as plain
///   struct fields or carry trait bounds that the precheck can't reason about
///   from the AST alone.
///
/// # Limitations
///
/// Generic arguments encoded inside a `PureType::Path` string (e.g. the `f32`
/// in `Box<f32>`) are **not** descended into. Doing so would require parsing
/// the type-string ourselves; that is left for a follow-up. The head walk
/// still pushes `Box`, and `Box<f32>: Hash` happens to be safe to skip because
/// `StdImplCache` advertises `Box: Hash` and the precheck does not yet inspect
/// the conditional bound.
fn walk_type_for_std_heads(ty: &PureType, std_impls: &StdImplCache, sink: &mut dyn FnMut(String)) {
    match ty {
        PureType::Path(p) => {
            let pre_generic = p.split('<').next().unwrap_or(p).trim();
            let head = pre_generic
                .rsplit("::")
                .next()
                .unwrap_or(pre_generic)
                .trim();
            if !head.is_empty()
                && (std_impls.is_primitive(head) || std_impls.is_std_container(head))
            {
                sink(head.to_string());
            }
        }
        PureType::Ref { ty, .. } => walk_type_for_std_heads(ty, std_impls, sink),
        PureType::Tuple(types) => {
            for t in types {
                walk_type_for_std_heads(t, std_impls, sink);
            }
        }
        PureType::Array { ty, .. } => walk_type_for_std_heads(ty, std_impls, sink),
        PureType::Slice(inner) => walk_type_for_std_heads(inner, std_impls, sink),
        // Fn / ImplTrait / TraitObject / Infer / Never / Other: nothing to
        // contribute to the precheck reject branch from the AST surface alone.
        _ => {}
    }
}

/// Walk a `PureType`, surfacing each **user-defined** (non-primitive,
/// non-std-container) head type name to `sink`.
///
/// This is the complement of [`walk_type_for_std_heads`]: where that function
/// only emits names accepted by `StdImplCache::is_primitive` /
/// `is_std_container`, this function emits names that are neither — i.e.
/// user-defined types that the typeflow path might have failed to resolve for
/// a given workspace.
///
/// The same structural recursion rules apply (Ref, Tuple, Array, Slice).
/// Generic arguments inside `PureType::Path` strings are not descended into
/// (same limitation as `walk_type_for_std_heads`).
///
/// Names starting with a lowercase letter that look like lifetimes or generic
/// parameters (single char, or all-lowercase short identifiers that are common
/// generic names) are conservatively skipped to avoid treating `T`, `K`, `V`,
/// `E`, etc. as user-defined struct names.
fn walk_type_for_user_heads(ty: &PureType, std_impls: &StdImplCache, sink: &mut dyn FnMut(String)) {
    match ty {
        PureType::Path(p) => {
            let pre_generic = p.split('<').next().unwrap_or(p).trim();
            let head = pre_generic
                .rsplit("::")
                .next()
                .unwrap_or(pre_generic)
                .trim();
            if head.is_empty() || std_impls.is_primitive(head) || std_impls.is_std_container(head) {
                return; // handled by walk_type_for_std_heads or skip
            }
            // Conservative filter: skip single-char identifiers and common
            // generic parameter names to avoid treating lifetime / type params
            // as user-defined struct names.
            if is_likely_generic_param(head) {
                return;
            }
            sink(head.to_string());
        }
        PureType::Ref { ty, .. } => walk_type_for_user_heads(ty, std_impls, sink),
        PureType::Tuple(types) => {
            for t in types {
                walk_type_for_user_heads(t, std_impls, sink);
            }
        }
        PureType::Array { ty, .. } => walk_type_for_user_heads(ty, std_impls, sink),
        PureType::Slice(inner) => walk_type_for_user_heads(inner, std_impls, sink),
        _ => {}
    }
}

/// Return true for identifiers that are almost certainly type/lifetime
/// parameters rather than user-defined struct / enum names.
///
/// Heuristic:
/// - Single ASCII letter (e.g. `T`, `K`, `V`, `E`, `F`, `N`).
/// - Two-char identifiers like `Ok`, `Err` are legitimate types and must NOT
///   be filtered; single-char is the primary filter.
/// - Common multi-char generic param names (`Self`, `Item`, `Output`,
///   `Error`) that appear frequently in trait bounds are kept so that a struct
///   literally named `Item` would still be captured — these are uncommon
///   enough that the FP risk is low.
fn is_likely_generic_param(name: &str) -> bool {
    let bytes = name.as_bytes();
    // Single ASCII letter.
    bytes.len() == 1 && bytes[0].is_ascii_alphabetic()
}

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

    /// Regression test (precheck-fix-1):
    /// `code_graph.children_of(struct_id)` returns every graph child, which
    /// for a unit struct used to include all methods on an `impl` block.
    /// Those methods' return types were then captured as "field types",
    /// causing `verify_precheck` to fabricate phantom fields (e.g.
    /// `[Result, Result, ...]`) for unit structs and to fail their derive
    /// check at baseline, rejecting every speculative mutation. The fix
    /// filters `children_of` to `SymbolKind::Field` only — this test locks
    /// down that filter.
    #[cfg(feature = "testing")]
    #[test]
    fn test_unit_struct_with_impl_methods_has_no_field_types() {
        use crate::testing::ContextBuilder;
        let source = r#"
#[derive(Debug, Clone, Default)]
pub struct UnitWithMethods;

impl UnitWithMethods {
    pub fn parse_a() -> Result<i32, String> { Ok(0) }
    pub fn parse_b() -> Result<i32, String> { Ok(0) }
}
"#;
        let ctx = ContextBuilder::new()
            .with_file("src/lib.rs", source)
            .build();

        // Find the struct's SymbolId
        let struct_id = ctx
            .registry
            .iter()
            .find(|(_, p)| p.name() == "UnitWithMethods")
            .map(|(id, _)| id)
            .expect("UnitWithMethods symbol should exist");

        // The struct has at least one derive recorded.
        let derives = ctx
            .derive_index
            .get_derives(struct_id)
            .expect("derives should be indexed");
        assert!(
            derives.iter().any(|d| d == "Default"),
            "Default should be among the derives, got {:?}",
            derives
        );

        // The struct is a unit struct → no real fields → field_type_names
        // must be empty (None) instead of leaking method return types.
        assert!(
            ctx.derive_index.get_field_types(struct_id).is_none(),
            "unit struct must not record any field types, got {:?}",
            ctx.derive_index.get_field_types(struct_id)
        );
    }
}