ryo_analysis/query/

typeflow_v2.rs

//! TypeFlowGraph V2 - Data-Oriented Design implementation
//!
//! High-performance type relationship tracking using:
//! - SoA (Structure of Arrays) for cache efficiency
//! - Inverted indices for O(1) impact analysis
//! - Span-free, String-free design following PureAST philosophy
//!
//! # Architecture
//!
//! ```text
//! TypeFlowGraphV2
//! ├── Data Storage (SoA)
//! │   ├── definitions: Vec<DefinitionData>
//! │   ├── usages: Vec<UsageData>
//! │   ├── generic_args: Vec<GenericArgData>
//! │   └── trait_bounds: Vec<TraitBoundData>
//! │
//! ├── Edge Relations (Side Tables)
//! │   ├── usage_to_args: Vec<SmallVec<[u32; 4]>>
//! │   └── node_to_bounds: FxHashMap<TypeNodeId, SmallVec<[u32; 2]>>
//! │
//! └── Lookup Tables (Inverted Index)
//!     ├── symbol_to_def: SecondaryMap<SymbolId, u32>
//!     ├── symbol_to_usages: FxHashMap<SymbolId, Vec<u32>>
//!     ├── trait_to_bounds: FxHashMap<SymbolId, Vec<u32>>
//!     └── container_fields: FxHashMap<SymbolId, Vec<SymbolId>>
//! ```

use crate::symbol::SymbolId;
use serde::{Deserialize, Serialize};
use slotmap::SecondaryMap;
use smallvec::SmallVec;
use std::collections::{HashMap, HashSet, VecDeque};

use super::graph_v2::{ChainDirection, ChainNode, ChainResult};

/// Reference kind for type usage.
///
/// Tracks ownership state: `T` vs `&T` vs `&mut T`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize)]
pub enum RefKind {
    /// Owned value: `T`
    #[default]
    Owned,
    /// Shared reference: `&T`
    Ref,
    /// Mutable reference: `&mut T`
    RefMut,
}

/// Kind of type definition.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum TypeDefKind {
    /// Struct definition.
    Struct,
    /// Enum definition.
    Enum,
    /// Trait definition.
    Trait,
    /// Type alias (type A = B).
    TypeAlias,
    /// Generic parameter definition (T in `fn foo<T>`).
    GenericParam,
    /// Associated type (type Output in trait).
    AssociatedType,
}

/// Context in which a type is used.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum UsageContext {
    /// Struct field type.
    FieldType,
    /// Function parameter type.
    ParamType,
    /// Return type.
    ReturnType,
    /// Local variable type annotation.
    LocalVar,
    /// Where clause or trait bound.
    WhereBound,
    /// Impl target (impl Foo).
    ImplTarget,
    /// Impl trait (impl Trait for Foo).
    ImplTrait,
    /// Use statement.
    UseStatement,
    /// Type cast (as Type).
    Cast,
    /// Turbofish (`::<Type>`).
    Turbofish,
    /// Enum variant field.
    VariantField,
    /// Const/static type.
    ConstType,
    /// Match expression scrutinee (the value being matched).
    /// Used to track where enums are matched for cascade analysis.
    MatchScrutinee,
    /// Associated function call: `Foo::new()`, `Foo::default()`, etc.
    /// Tracks the type referenced by the path prefix of the call.
    ConstructorCall,
    /// Struct literal: `Foo { field: value }`.
    /// Tracks the type being instantiated.
    StructLiteral,
}

// ============================================================================
// Node ID Types
// ============================================================================

/// Unified node ID with kind encoded in top 2 bits.
/// Allows a single ID type to reference any node kind.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Serialize, Deserialize)]
#[repr(transparent)]
pub struct TypeNodeId(u32);

impl TypeNodeId {
    const KIND_SHIFT: u32 = 30;
    const KIND_MASK: u32 = 0xC000_0000;
    const INDEX_MASK: u32 = 0x3FFF_FFFF;

    const KIND_DEF: u32 = 0;
    const KIND_USAGE: u32 = 1;
    const KIND_ARG: u32 = 2;
    const KIND_BOUND: u32 = 3;

    /// Create a definition node ID.
    #[inline]
    pub const fn def(idx: u32) -> Self {
        Self((Self::KIND_DEF << Self::KIND_SHIFT) | idx)
    }

    /// Create a usage node ID.
    #[inline]
    pub const fn usage(idx: u32) -> Self {
        Self((Self::KIND_USAGE << Self::KIND_SHIFT) | idx)
    }

    /// Create a generic arg node ID.
    #[inline]
    pub const fn arg(idx: u32) -> Self {
        Self((Self::KIND_ARG << Self::KIND_SHIFT) | idx)
    }

    /// Create a trait bound node ID.
    #[inline]
    pub const fn bound(idx: u32) -> Self {
        Self((Self::KIND_BOUND << Self::KIND_SHIFT) | idx)
    }

    /// Get the node kind.
    #[inline]
    pub const fn kind(self) -> NodeKind {
        match (self.0 & Self::KIND_MASK) >> Self::KIND_SHIFT {
            Self::KIND_DEF => NodeKind::Definition,
            Self::KIND_USAGE => NodeKind::Usage,
            Self::KIND_ARG => NodeKind::GenericArg,
            Self::KIND_BOUND => NodeKind::TraitBound,
            _ => unreachable!(),
        }
    }

    /// Get the raw index within the kind's array.
    #[inline]
    pub const fn index(self) -> u32 {
        self.0 & Self::INDEX_MASK
    }

    /// Get raw u32 value.
    #[inline]
    pub const fn as_u32(self) -> u32 {
        self.0
    }
}

/// Node kind discriminant.
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub enum NodeKind {
    /// Type definition site (where the type is declared).
    Definition,
    /// Concrete usage of a type at a value position.
    Usage,
    /// Generic argument position (e.g. `T` in `Vec<T>`).
    GenericArg,
    /// Trait-bound position (e.g. `T: Trait`).
    TraitBound,
}

// ============================================================================
// Data Structures (SoA Components)
// ============================================================================

/// Definition node data (16 bytes).
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
pub struct DefinitionData {
    /// Symbol ID of the defined type.
    pub symbol_id: SymbolId,
    /// Kind of type definition.
    pub kind: TypeDefKind,
}

/// Usage node data - Span-free, String-free.
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
pub struct UsageData {
    /// Resolved type definition (if available).
    pub resolved: Option<SymbolId>,
    /// Context in which the type is used.
    pub context: UsageContext,
    /// Reference kind (owned, ref, ref mut).
    pub ref_kind: RefKind,
    /// Container symbol that owns this usage (e.g., the function or struct).
    ///
    /// For `UsageContext::ParamType` / `ReturnType`, this is the function's SymbolId.
    /// For `UsageContext::FieldType`, this is the struct/enum's SymbolId.
    pub container: Option<SymbolId>,
}

/// Generic argument node data (16 bytes).
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
pub struct GenericArgData {
    /// Parent usage's index.
    pub parent_usage: u32,
    /// Argument index (0, 1, 2...).
    pub arg_index: u8,
    /// Resolved type (if concrete).
    pub resolved: Option<SymbolId>,
}

/// Trait bound node data (16 bytes).
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
pub struct TraitBoundData {
    /// Target node being bounded.
    pub target: TypeNodeId,
    /// Trait's SymbolId (if resolved).
    pub trait_id: Option<SymbolId>,
}

// ============================================================================
// Lookup Table (Inverted Index)
// ============================================================================

/// Inverted indices for O(1) lookups.
#[derive(Clone, Default, Debug, Serialize, Deserialize)]
pub struct LookupTable {
    /// SymbolId → definition index (1:1).
    pub symbol_to_def: SecondaryMap<SymbolId, u32>,

    /// Resolved type SymbolId → usage indices (1:N).
    /// Reverse lookup: "what usages reference this type?"
    pub symbol_to_usages: HashMap<SymbolId, Vec<u32>>,

    /// Container SymbolId → usage indices (1:N).
    /// Forward lookup: "what types does this container use?"
    pub container_to_usages: HashMap<SymbolId, Vec<u32>>,

    /// Trait SymbolId → bound indices (1:N).
    pub trait_to_bounds: HashMap<SymbolId, Vec<u32>>,

    /// Container SymbolId → field type SymbolIds (Contains relation).
    pub container_fields: HashMap<SymbolId, Vec<SymbolId>>,
}

// ============================================================================
// TypeFlowGraph V2
// ============================================================================

/// Type relationship graph (V2 - Data-Oriented Design).
///
/// NOTE: Deserialize is NOT derived because TypeFlowGraphV2 contains
/// SecondaryMap<SymbolId, ...> and HashMap<SymbolId, ...>.
/// SymbolId is process-specific. Serialize is kept for debugging/inspection.
#[derive(Clone, Default, Debug, Serialize)]
pub struct TypeFlowGraphV2 {
    // === Data Storage (SoA) ===
    /// Definition nodes.
    pub definitions: Vec<DefinitionData>,

    /// Usage nodes.
    pub usages: Vec<UsageData>,

    /// Generic argument nodes.
    pub generic_args: Vec<GenericArgData>,

    /// Trait bound nodes.
    pub trait_bounds: Vec<TraitBoundData>,

    // === Edge Relations (Side Tables) ===
    /// Usage → generic args (1:N).
    pub usage_to_args: Vec<SmallVec<[u32; 4]>>,

    /// Node → trait bounds (1:N).
    pub node_to_bounds: HashMap<TypeNodeId, SmallVec<[u32; 2]>>,

    // === Lookup Tables ===
    /// Reverse-lookup tables (`SymbolId` → node indices, etc.) used for
    /// fast querying.
    pub lookup: LookupTable,
}

/// Impact of changing a type.
#[derive(Debug, Clone, Default)]
pub struct TypeImpactV2 {
    /// Direct usage indices.
    pub direct_usages: Vec<u32>,
    /// Trait bound indices.
    pub bound_usages: Vec<u32>,
    /// Types that contain this type as a field.
    pub containing_types: Vec<SymbolId>,
    /// Generic usage indices.
    pub generic_usages: Vec<u32>,
}

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

    /// Create with pre-allocated capacity.
    pub fn with_capacity(defs: usize, usages: usize) -> Self {
        Self {
            definitions: Vec::with_capacity(defs),
            usages: Vec::with_capacity(usages),
            generic_args: Vec::with_capacity(usages), // Usually similar count
            trait_bounds: Vec::new(),
            usage_to_args: Vec::with_capacity(usages),
            node_to_bounds: HashMap::new(),
            lookup: LookupTable::default(),
        }
    }

    // ========================================================================
    // Node Management
    // ========================================================================

    /// Add a type definition node.
    pub fn add_definition(&mut self, symbol_id: SymbolId, kind: TypeDefKind) -> TypeNodeId {
        // Return existing if already registered
        if let Some(&idx) = self.lookup.symbol_to_def.get(symbol_id) {
            return TypeNodeId::def(idx);
        }

        let idx = self.definitions.len() as u32;
        self.definitions.push(DefinitionData { symbol_id, kind });
        self.lookup.symbol_to_def.insert(symbol_id, idx);

        TypeNodeId::def(idx)
    }

    /// Add a type usage node.
    pub fn add_usage(
        &mut self,
        context: UsageContext,
        ref_kind: RefKind,
        resolved: Option<SymbolId>,
        container: Option<SymbolId>,
    ) -> TypeNodeId {
        let idx = self.usages.len() as u32;
        self.usages.push(UsageData {
            resolved,
            context,
            ref_kind,
            container,
        });

        // Initialize empty args list
        self.usage_to_args.push(SmallVec::new());

        // Register in reverse lookup (type → usages)
        if let Some(symbol_id) = resolved {
            self.lookup
                .symbol_to_usages
                .entry(symbol_id)
                .or_default()
                .push(idx);
        }

        // Register in forward lookup (container → usages)
        if let Some(cid) = container {
            self.lookup
                .container_to_usages
                .entry(cid)
                .or_default()
                .push(idx);
        }

        TypeNodeId::usage(idx)
    }

    /// Add a generic argument node.
    pub fn add_generic_arg(
        &mut self,
        parent: TypeNodeId,
        arg_index: u8,
        resolved: Option<SymbolId>,
    ) -> TypeNodeId {
        let idx = self.generic_args.len() as u32;
        self.generic_args.push(GenericArgData {
            parent_usage: parent.index(),
            arg_index,
            resolved,
        });

        // Link to parent usage
        if parent.kind() == NodeKind::Usage {
            if let Some(args) = self.usage_to_args.get_mut(parent.index() as usize) {
                args.push(idx);
            }

            // Register in reverse lookup (type → usages) so type_users() can find
            // containers that reference this type via generic arguments (e.g., Vec<Foo>).
            if let Some(symbol_id) = resolved {
                self.lookup
                    .symbol_to_usages
                    .entry(symbol_id)
                    .or_default()
                    .push(parent.index());
            }
        }

        TypeNodeId::arg(idx)
    }

    /// Add a trait bound node.
    pub fn add_trait_bound(
        &mut self,
        target: TypeNodeId,
        trait_id: Option<SymbolId>,
    ) -> TypeNodeId {
        let idx = self.trait_bounds.len() as u32;
        self.trait_bounds.push(TraitBoundData { target, trait_id });

        // Link to target node
        self.node_to_bounds.entry(target).or_default().push(idx);

        // Register in lookup table
        if let Some(tid) = trait_id {
            self.lookup
                .trait_to_bounds
                .entry(tid)
                .or_default()
                .push(idx);
        }

        TypeNodeId::bound(idx)
    }

    /// Register a container-field relationship.
    pub fn add_contains(&mut self, container: SymbolId, field_type: SymbolId) {
        self.lookup
            .container_fields
            .entry(container)
            .or_default()
            .push(field_type);
    }

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

    /// Get the definition node for a symbol (O(1)).
    #[inline]
    pub fn definition(&self, id: SymbolId) -> Option<TypeNodeId> {
        self.lookup
            .symbol_to_def
            .get(id)
            .map(|&idx| TypeNodeId::def(idx))
    }

    /// Get all usage indices for a type (O(1)).
    #[inline]
    pub fn usages(&self, id: SymbolId) -> impl Iterator<Item = TypeNodeId> + '_ {
        self.lookup
            .symbol_to_usages
            .get(&id)
            .into_iter()
            .flat_map(|v| v.iter().map(|&idx| TypeNodeId::usage(idx)))
    }

    /// Get the number of usages for a type (O(1)).
    #[inline]
    pub fn usage_count(&self, id: SymbolId) -> usize {
        self.lookup.symbol_to_usages.get(&id).map_or(0, |v| v.len())
    }

    /// Forward lookup: what types does this container use? (O(1) + result count)
    ///
    /// Returns resolved SymbolIds of types referenced by the container.
    /// Includes both directly resolved types and types referenced via generic arguments
    /// (e.g., `Foo` in `Vec<Foo>`).
    pub fn types_used_by(&self, container: SymbolId) -> impl Iterator<Item = SymbolId> + '_ {
        self.lookup
            .container_to_usages
            .get(&container)
            .into_iter()
            .flat_map(|indices| {
                indices.iter().flat_map(|&idx| {
                    let mut types = SmallVec::<[SymbolId; 4]>::new();
                    // Direct type reference (e.g., field: Foo)
                    if let Some(u) = self.usages.get(idx as usize) {
                        if let Some(resolved) = u.resolved {
                            types.push(resolved);
                        }
                    }
                    // Generic argument types (e.g., Foo in Vec<Foo>)
                    if let Some(args) = self.usage_to_args.get(idx as usize) {
                        for &arg_idx in args.iter() {
                            if let Some(arg) = self.generic_args.get(arg_idx as usize) {
                                if let Some(resolved) = arg.resolved {
                                    types.push(resolved);
                                }
                            }
                        }
                    }
                    types.into_iter()
                })
            })
    }

    /// Reverse lookup: what containers use this type? (O(1) + result count)
    ///
    /// Returns SymbolIds of containers that reference the given type.
    pub fn type_users(&self, type_id: SymbolId) -> impl Iterator<Item = SymbolId> + '_ {
        self.lookup
            .symbol_to_usages
            .get(&type_id)
            .into_iter()
            .flat_map(|indices| {
                indices
                    .iter()
                    .filter_map(|&idx| self.usages.get(idx as usize).and_then(|u| u.container))
            })
    }

    // ========================================================================
    // Chain traversal (BFS)
    // ========================================================================

    /// BFS traversal: find all containers that use this type, transitively.
    ///
    /// "Type A is used by container B; B (as a type) is used by container C" etc.
    pub fn type_users_chain(&self, start: SymbolId, max_depth: usize) -> Vec<ChainNode> {
        self.traverse_type_chain(start, max_depth, ChainDirection::TypeUsers)
    }

    /// BFS traversal: find all types used by this container, transitively.
    ///
    /// "Container A uses type B; B (as a container) uses type C" etc.
    pub fn types_used_by_chain(&self, start: SymbolId, max_depth: usize) -> Vec<ChainNode> {
        self.traverse_type_chain(start, max_depth, ChainDirection::TypeDeps)
    }

    /// Full chain analysis with statistics.
    ///
    /// Returns `ChainResult` compatible with CodeGraphV2's `analyze_chain()`.
    pub fn analyze_type_chain(
        &self,
        start: SymbolId,
        max_depth: usize,
        direction: ChainDirection,
    ) -> ChainResult {
        let nodes = self.traverse_type_chain(start, max_depth, direction);

        let mut by_depth: HashMap<usize, usize> = HashMap::new();
        for node in &nodes {
            *by_depth.entry(node.depth).or_default() += 1;
        }

        let max_actual_depth = nodes.iter().map(|n| n.depth).max().unwrap_or(0);

        ChainResult {
            start,
            direction,
            max_depth,
            nodes,
            max_actual_depth,
            by_depth,
        }
    }

    /// Internal BFS implementation for type chain traversal.
    fn traverse_type_chain(
        &self,
        start: SymbolId,
        max_depth: usize,
        direction: ChainDirection,
    ) -> Vec<ChainNode> {
        let mut result = Vec::new();
        let mut visited = HashSet::new();
        let mut queue = VecDeque::new();

        visited.insert(start);
        queue.push_back((start, 0usize));

        while let Some((current, depth)) = queue.pop_front() {
            if depth > 0 {
                result.push(ChainNode {
                    symbol: current,
                    depth,
                });
            }

            if depth >= max_depth {
                continue;
            }

            let neighbors: Vec<SymbolId> = match direction {
                ChainDirection::TypeUsers => self.type_users(current).collect(),
                ChainDirection::TypeDeps => self.types_used_by(current).collect(),
                ChainDirection::Callers | ChainDirection::Callees => {
                    unreachable!("Callers/Callees must use CodeGraphV2")
                }
            };

            for neighbor in neighbors {
                if !visited.contains(&neighbor) {
                    visited.insert(neighbor);
                    queue.push_back((neighbor, depth + 1));
                }
            }
        }

        result
    }

    /// Get generic arguments for a usage node (O(1)).
    #[inline]
    pub fn generic_args(&self, node: TypeNodeId) -> impl Iterator<Item = TypeNodeId> + '_ {
        let args = if node.kind() == NodeKind::Usage {
            self.usage_to_args.get(node.index() as usize)
        } else {
            None
        };
        args.into_iter()
            .flat_map(|v| v.iter().map(|&idx| TypeNodeId::arg(idx)))
    }

    /// Get trait bounds for a node (O(1)).
    #[inline]
    pub fn trait_bounds(&self, node: TypeNodeId) -> impl Iterator<Item = TypeNodeId> + '_ {
        self.node_to_bounds
            .get(&node)
            .into_iter()
            .flat_map(|v| v.iter().map(|&idx| TypeNodeId::bound(idx)))
    }

    /// Get definition data by ID.
    #[inline]
    pub fn get_definition(&self, id: TypeNodeId) -> Option<&DefinitionData> {
        if id.kind() == NodeKind::Definition {
            self.definitions.get(id.index() as usize)
        } else {
            None
        }
    }

    /// Get usage data by ID.
    #[inline]
    pub fn get_usage(&self, id: TypeNodeId) -> Option<&UsageData> {
        if id.kind() == NodeKind::Usage {
            self.usages.get(id.index() as usize)
        } else {
            None
        }
    }

    /// Get generic arg data by ID.
    #[inline]
    pub fn get_generic_arg(&self, id: TypeNodeId) -> Option<&GenericArgData> {
        if id.kind() == NodeKind::GenericArg {
            self.generic_args.get(id.index() as usize)
        } else {
            None
        }
    }

    /// Get trait bound data by ID.
    #[inline]
    pub fn get_trait_bound(&self, id: TypeNodeId) -> Option<&TraitBoundData> {
        if id.kind() == NodeKind::TraitBound {
            self.trait_bounds.get(id.index() as usize)
        } else {
            None
        }
    }

    // ========================================================================
    // Impact Analysis
    // ========================================================================

    /// Analyze the impact of changing a type (O(1) + result count).
    pub fn impact(&self, id: SymbolId) -> TypeImpactV2 {
        TypeImpactV2 {
            // Direct usages: O(1)
            direct_usages: self
                .lookup
                .symbol_to_usages
                .get(&id)
                .cloned()
                .unwrap_or_default(),

            // Trait bounds: O(1)
            bound_usages: self
                .lookup
                .trait_to_bounds
                .get(&id)
                .cloned()
                .unwrap_or_default(),

            // Containing types: O(containers)
            containing_types: self
                .lookup
                .container_fields
                .iter()
                .filter(|(_, fields)| fields.contains(&id))
                .map(|(&container, _)| container)
                .collect(),

            // Generic usages: computed separately if needed
            generic_usages: vec![],
        }
    }

    // ========================================================================
    // Statistics
    // ========================================================================

    /// Total number of nodes.
    pub fn node_count(&self) -> usize {
        self.definitions.len()
            + self.usages.len()
            + self.generic_args.len()
            + self.trait_bounds.len()
    }

    /// Number of definitions.
    pub fn definition_count(&self) -> usize {
        self.definitions.len()
    }

    /// Number of usages.
    pub fn usages_count(&self) -> usize {
        self.usages.len()
    }

    /// Check if empty.
    pub fn is_empty(&self) -> bool {
        self.definitions.is_empty() && self.usages.is_empty()
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::symbol::{SymbolPath, SymbolRegistry};
    use crate::SymbolKind;

    fn create_test_registry() -> SymbolRegistry {
        let mut registry = SymbolRegistry::new();
        let _ = registry.register(SymbolPath::parse("test::Foo").unwrap(), SymbolKind::Struct);
        let _ = registry.register(SymbolPath::parse("test::Bar").unwrap(), SymbolKind::Struct);
        let _ = registry.register(SymbolPath::parse("test::Clone").unwrap(), SymbolKind::Trait);
        registry
    }

    fn lookup(registry: &SymbolRegistry, path: &str) -> SymbolId {
        let symbol_path = SymbolPath::parse(path).unwrap();
        registry.lookup(&symbol_path).unwrap()
    }

    #[test]
    fn test_new_graph() {
        let graph = TypeFlowGraphV2::new();
        assert!(graph.is_empty());
        assert_eq!(graph.node_count(), 0);
    }

    #[test]
    fn test_add_definition() {
        let mut graph = TypeFlowGraphV2::new();
        let registry = create_test_registry();

        let foo_id = lookup(&registry, "test::Foo");
        let node = graph.add_definition(foo_id, TypeDefKind::Struct);

        assert_eq!(graph.definition_count(), 1);
        assert_eq!(graph.definition(foo_id), Some(node));
        assert_eq!(node.kind(), NodeKind::Definition);

        // Adding same definition returns existing node
        let node2 = graph.add_definition(foo_id, TypeDefKind::Struct);
        assert_eq!(node, node2);
        assert_eq!(graph.definition_count(), 1);
    }

    #[test]
    fn test_add_usage() {
        let mut graph = TypeFlowGraphV2::new();
        let registry = create_test_registry();

        let foo_id = lookup(&registry, "test::Foo");
        graph.add_definition(foo_id, TypeDefKind::Struct);

        let usage = graph.add_usage(UsageContext::FieldType, RefKind::Owned, Some(foo_id), None);

        assert_eq!(graph.usage_count(foo_id), 1);
        assert_eq!(usage.kind(), NodeKind::Usage);

        let usages: Vec<_> = graph.usages(foo_id).collect();
        assert_eq!(usages.len(), 1);
        assert_eq!(usages[0], usage);
    }

    #[test]
    fn test_add_usage_with_container() {
        let mut graph = TypeFlowGraphV2::new();
        let registry = create_test_registry();

        let foo_id = lookup(&registry, "test::Foo");
        let bar_id = lookup(&registry, "test::Bar");
        graph.add_definition(foo_id, TypeDefKind::Struct);
        graph.add_definition(bar_id, TypeDefKind::Struct);

        // Bar has a field of type Foo
        let usage = graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            Some(foo_id),
            Some(bar_id),
        );

        // Verify container is stored
        let data = graph.get_usage(usage).unwrap();
        assert_eq!(data.container, Some(bar_id));
        assert_eq!(data.resolved, Some(foo_id));
    }

    #[test]
    fn test_container_to_usages_index() {
        let mut graph = TypeFlowGraphV2::new();
        let registry = create_test_registry();

        let foo_id = lookup(&registry, "test::Foo");
        let bar_id = lookup(&registry, "test::Bar");
        let clone_id = lookup(&registry, "test::Clone");
        graph.add_definition(foo_id, TypeDefKind::Struct);
        graph.add_definition(bar_id, TypeDefKind::Struct);
        graph.add_definition(clone_id, TypeDefKind::Trait);

        // Bar uses Foo (field) and Clone (trait bound) — two usages from same container
        graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            Some(foo_id),
            Some(bar_id),
        );
        graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            Some(clone_id),
            Some(bar_id),
        );

        // container_to_usages should have 2 entries for bar_id
        let indices = graph.lookup.container_to_usages.get(&bar_id).unwrap();
        assert_eq!(indices.len(), 2);
    }

    #[test]
    fn test_types_used_by() {
        let mut graph = TypeFlowGraphV2::new();
        let registry = create_test_registry();

        let foo_id = lookup(&registry, "test::Foo");
        let bar_id = lookup(&registry, "test::Bar");
        let clone_id = lookup(&registry, "test::Clone");
        graph.add_definition(foo_id, TypeDefKind::Struct);
        graph.add_definition(bar_id, TypeDefKind::Struct);
        graph.add_definition(clone_id, TypeDefKind::Trait);

        // Bar uses Foo and Clone
        graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            Some(foo_id),
            Some(bar_id),
        );
        graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            Some(clone_id),
            Some(bar_id),
        );

        let used: Vec<_> = graph.types_used_by(bar_id).collect();
        assert_eq!(used.len(), 2);
        assert!(used.contains(&foo_id));
        assert!(used.contains(&clone_id));

        // Foo uses nothing
        let used_by_foo: Vec<_> = graph.types_used_by(foo_id).collect();
        assert!(used_by_foo.is_empty());
    }

    /// Reproducer: types_used_by should return types referenced via generic wrappers.
    ///
    /// Simulates: `struct Bar { items: Vec<Foo> }`
    /// The builder produces:
    ///   add_usage(FieldType, Owned, resolved=None, container=Some(bar_id))  ← Vec unresolved
    ///   add_generic_arg(parent_usage, 0, resolved=Some(foo_id))             ← Foo as arg
    ///
    /// Expected: types_used_by(bar_id) should include foo_id.
    #[test]
    fn test_types_used_by_includes_generic_arg_types() {
        let mut graph = TypeFlowGraphV2::new();
        let registry = create_test_registry();

        let foo_id = lookup(&registry, "test::Foo");
        let bar_id = lookup(&registry, "test::Bar");
        graph.add_definition(foo_id, TypeDefKind::Struct);
        graph.add_definition(bar_id, TypeDefKind::Struct);

        // Bar has field: Vec<Foo>
        // Vec is unresolved (std type) → resolved=None
        let vec_usage = graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            None, // Vec can't be resolved
            Some(bar_id),
        );

        // Foo is the generic argument of Vec
        graph.add_generic_arg(vec_usage, 0, Some(foo_id));

        // types_used_by(bar_id) should still find Foo via the generic arg
        let used: Vec<_> = graph.types_used_by(bar_id).collect();
        assert!(
            used.contains(&foo_id),
            "types_used_by should include Foo referenced via Vec<Foo>, got: {:?}",
            used
        );
    }

    #[test]
    fn test_type_users() {
        let mut graph = TypeFlowGraphV2::new();
        let registry = create_test_registry();

        let foo_id = lookup(&registry, "test::Foo");
        let bar_id = lookup(&registry, "test::Bar");
        let clone_id = lookup(&registry, "test::Clone");
        graph.add_definition(foo_id, TypeDefKind::Struct);
        graph.add_definition(bar_id, TypeDefKind::Struct);
        graph.add_definition(clone_id, TypeDefKind::Trait);

        // Bar and Clone both use Foo
        graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            Some(foo_id),
            Some(bar_id),
        );
        graph.add_usage(
            UsageContext::ParamType,
            RefKind::Ref,
            Some(foo_id),
            Some(clone_id),
        );

        let users: Vec<_> = graph.type_users(foo_id).collect();
        assert_eq!(users.len(), 2);
        assert!(users.contains(&bar_id));
        assert!(users.contains(&clone_id));

        // Nobody uses Clone
        let users_of_clone: Vec<_> = graph.type_users(clone_id).collect();
        assert!(users_of_clone.is_empty());
    }

    /// Reproducer: type_users should return containers that reference a type via generic args.
    ///
    /// Simulates: `struct Bar { items: Vec<Foo> }`
    /// Expected: type_users(foo_id) should include bar_id.
    #[test]
    fn test_type_users_includes_generic_arg_references() {
        let mut graph = TypeFlowGraphV2::new();
        let registry = create_test_registry();

        let foo_id = lookup(&registry, "test::Foo");
        let bar_id = lookup(&registry, "test::Bar");
        graph.add_definition(foo_id, TypeDefKind::Struct);
        graph.add_definition(bar_id, TypeDefKind::Struct);

        // Bar has field: Vec<Foo>
        // Vec is unresolved → resolved=None
        let vec_usage =
            graph.add_usage(UsageContext::FieldType, RefKind::Owned, None, Some(bar_id));

        // Foo is the generic argument of Vec
        graph.add_generic_arg(vec_usage, 0, Some(foo_id));

        // type_users(foo_id) should find Bar as a user via the generic arg
        let users: Vec<_> = graph.type_users(foo_id).collect();
        assert!(
            users.contains(&bar_id),
            "type_users should include Bar which references Foo via Vec<Foo>, got: {:?}",
            users
        );
    }

    #[test]
    fn test_type_users_without_container() {
        let mut graph = TypeFlowGraphV2::new();
        let registry = create_test_registry();

        let foo_id = lookup(&registry, "test::Foo");
        graph.add_definition(foo_id, TypeDefKind::Struct);

        // Usage without container (legacy behavior)
        graph.add_usage(UsageContext::FieldType, RefKind::Owned, Some(foo_id), None);

        // type_users should return empty since no container is set
        let users: Vec<_> = graph.type_users(foo_id).collect();
        assert!(users.is_empty());

        // But usages() still works (returns TypeNodeId)
        assert_eq!(graph.usage_count(foo_id), 1);
    }

    #[test]
    fn test_generic_args() {
        let mut graph = TypeFlowGraphV2::new();
        let registry = create_test_registry();

        let foo_id = lookup(&registry, "test::Foo");
        graph.add_definition(foo_id, TypeDefKind::Struct);

        // Vec<Foo>
        let vec_usage = graph.add_usage(UsageContext::FieldType, RefKind::Owned, None, None);
        let arg = graph.add_generic_arg(vec_usage, 0, Some(foo_id));

        let args: Vec<_> = graph.generic_args(vec_usage).collect();
        assert_eq!(args.len(), 1);
        assert_eq!(args[0], arg);
        assert_eq!(arg.kind(), NodeKind::GenericArg);
    }

    #[test]
    fn test_trait_bounds() {
        let mut graph = TypeFlowGraphV2::new();
        let registry = create_test_registry();

        let clone_id = lookup(&registry, "test::Clone");
        graph.add_definition(clone_id, TypeDefKind::Trait);

        // T: Clone
        let param = graph.add_usage(UsageContext::WhereBound, RefKind::Owned, None, None);
        let bound = graph.add_trait_bound(param, Some(clone_id));

        let bounds: Vec<_> = graph.trait_bounds(param).collect();
        assert_eq!(bounds.len(), 1);
        assert_eq!(bounds[0], bound);
        assert_eq!(bound.kind(), NodeKind::TraitBound);
    }

    #[test]
    fn test_impact_analysis() {
        let mut graph = TypeFlowGraphV2::new();
        let registry = create_test_registry();

        let foo_id = lookup(&registry, "test::Foo");
        let bar_id = lookup(&registry, "test::Bar");

        graph.add_definition(foo_id, TypeDefKind::Struct);
        graph.add_definition(bar_id, TypeDefKind::Struct);

        // Bar has a field of type Foo
        graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            Some(foo_id),
            Some(bar_id),
        );
        graph.add_contains(bar_id, foo_id);

        let impact = graph.impact(foo_id);
        assert_eq!(impact.direct_usages.len(), 1);
        assert!(impact.containing_types.contains(&bar_id));
    }

    #[test]
    fn test_type_node_id_encoding() {
        // Test that different kinds produce different IDs
        let def = TypeNodeId::def(100);
        let usage = TypeNodeId::usage(100);
        let arg = TypeNodeId::arg(100);
        let bound = TypeNodeId::bound(100);

        assert_ne!(def, usage);
        assert_ne!(usage, arg);
        assert_ne!(arg, bound);

        // Test round-trip
        assert_eq!(def.kind(), NodeKind::Definition);
        assert_eq!(def.index(), 100);

        assert_eq!(usage.kind(), NodeKind::Usage);
        assert_eq!(usage.index(), 100);

        assert_eq!(arg.kind(), NodeKind::GenericArg);
        assert_eq!(arg.index(), 100);

        assert_eq!(bound.kind(), NodeKind::TraitBound);
        assert_eq!(bound.index(), 100);
    }

    // ====================================================================
    // BFS chain traversal tests
    // ====================================================================

    /// Build a type graph:
    /// fn process(x: Foo) -> Bar   (process uses Foo, Bar)
    /// fn render(y: Bar)           (render uses Bar)
    /// struct Baz { f: Foo }       (Baz uses Foo)
    fn build_chain_test_graph() -> (
        TypeFlowGraphV2,
        SymbolId,
        SymbolId,
        SymbolId,
        SymbolId,
        SymbolId,
    ) {
        let mut registry = SymbolRegistry::new();
        let foo_id = registry
            .register(SymbolPath::parse("test::Foo").unwrap(), SymbolKind::Struct)
            .unwrap();
        let bar_id = registry
            .register(SymbolPath::parse("test::Bar").unwrap(), SymbolKind::Struct)
            .unwrap();
        let baz_id = registry
            .register(SymbolPath::parse("test::Baz").unwrap(), SymbolKind::Struct)
            .unwrap();
        let process_id = registry
            .register(
                SymbolPath::parse("test::process").unwrap(),
                SymbolKind::Function,
            )
            .unwrap();
        let render_id = registry
            .register(
                SymbolPath::parse("test::render").unwrap(),
                SymbolKind::Function,
            )
            .unwrap();

        let mut graph = TypeFlowGraphV2::new();
        graph.add_definition(foo_id, TypeDefKind::Struct);
        graph.add_definition(bar_id, TypeDefKind::Struct);
        graph.add_definition(baz_id, TypeDefKind::Struct);

        // process uses Foo (param) and Bar (return)
        graph.add_usage(
            UsageContext::ParamType,
            RefKind::Owned,
            Some(foo_id),
            Some(process_id),
        );
        graph.add_usage(
            UsageContext::ReturnType,
            RefKind::Owned,
            Some(bar_id),
            Some(process_id),
        );

        // render uses Bar (param)
        graph.add_usage(
            UsageContext::ParamType,
            RefKind::Owned,
            Some(bar_id),
            Some(render_id),
        );

        // Baz uses Foo (field)
        graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            Some(foo_id),
            Some(baz_id),
        );

        (graph, foo_id, bar_id, baz_id, process_id, render_id)
    }

    #[test]
    fn test_type_users_chain_depth1() {
        let (graph, foo_id, _, baz_id, process_id, _) = build_chain_test_graph();

        // Foo is used by process (param) and Baz (field)
        let chain = graph.type_users_chain(foo_id, 1);
        let symbols: Vec<_> = chain.iter().map(|n| n.symbol).collect();
        assert_eq!(chain.len(), 2);
        assert!(symbols.contains(&process_id));
        assert!(symbols.contains(&baz_id));
        assert!(chain.iter().all(|n| n.depth == 1));
    }

    #[test]
    fn test_types_used_by_chain_depth1() {
        let (graph, foo_id, bar_id, _, process_id, _) = build_chain_test_graph();

        // process uses Foo and Bar
        let chain = graph.types_used_by_chain(process_id, 1);
        let symbols: Vec<_> = chain.iter().map(|n| n.symbol).collect();
        assert_eq!(chain.len(), 2);
        assert!(symbols.contains(&foo_id));
        assert!(symbols.contains(&bar_id));
    }

    #[test]
    fn test_type_users_chain_transitive() {
        // Build a chain: StructA is used by StructB (field), StructB is used by fn_c (param)
        let mut registry = SymbolRegistry::new();
        let a_id = registry
            .register(SymbolPath::parse("test::A").unwrap(), SymbolKind::Struct)
            .unwrap();
        let b_id = registry
            .register(SymbolPath::parse("test::B").unwrap(), SymbolKind::Struct)
            .unwrap();
        let c_id = registry
            .register(SymbolPath::parse("test::c").unwrap(), SymbolKind::Function)
            .unwrap();

        let mut graph = TypeFlowGraphV2::new();
        graph.add_definition(a_id, TypeDefKind::Struct);
        graph.add_definition(b_id, TypeDefKind::Struct);

        // B uses A (field)
        graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            Some(a_id),
            Some(b_id),
        );
        // c uses B (param)
        graph.add_usage(
            UsageContext::ParamType,
            RefKind::Owned,
            Some(b_id),
            Some(c_id),
        );

        // type_users_chain(A, 2): A ← B (depth 1), B ← c (depth 2)
        let chain = graph.type_users_chain(a_id, 2);
        assert_eq!(chain.len(), 2);
        assert_eq!(chain[0].symbol, b_id);
        assert_eq!(chain[0].depth, 1);
        assert_eq!(chain[1].symbol, c_id);
        assert_eq!(chain[1].depth, 2);
    }

    #[test]
    fn test_type_users_chain_respects_max_depth() {
        let mut registry = SymbolRegistry::new();
        let a_id = registry
            .register(SymbolPath::parse("test::A").unwrap(), SymbolKind::Struct)
            .unwrap();
        let b_id = registry
            .register(SymbolPath::parse("test::B").unwrap(), SymbolKind::Struct)
            .unwrap();
        let c_id = registry
            .register(SymbolPath::parse("test::c").unwrap(), SymbolKind::Function)
            .unwrap();

        let mut graph = TypeFlowGraphV2::new();
        graph.add_definition(a_id, TypeDefKind::Struct);
        graph.add_definition(b_id, TypeDefKind::Struct);

        graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            Some(a_id),
            Some(b_id),
        );
        graph.add_usage(
            UsageContext::ParamType,
            RefKind::Owned,
            Some(b_id),
            Some(c_id),
        );

        // max_depth=1: should only find B, not c
        let chain = graph.type_users_chain(a_id, 1);
        assert_eq!(chain.len(), 1);
        assert_eq!(chain[0].symbol, b_id);
        assert_eq!(chain[0].depth, 1);
    }

    #[test]
    fn test_analyze_type_chain_statistics() {
        let (graph, foo_id, _, _, _, _) = build_chain_test_graph();

        let result = graph.analyze_type_chain(foo_id, 5, ChainDirection::TypeUsers);
        assert_eq!(result.start, foo_id);
        assert_eq!(result.direction, ChainDirection::TypeUsers);
        assert_eq!(result.max_depth, 5);
        assert_eq!(result.nodes.len(), 2); // process and Baz
        assert_eq!(result.max_actual_depth, 1);
        assert_eq!(*result.by_depth.get(&1).unwrap_or(&0), 2);
    }

    #[test]
    fn test_chain_no_cycles() {
        // Ensure BFS doesn't loop on circular references
        let mut registry = SymbolRegistry::new();
        let a_id = registry
            .register(SymbolPath::parse("test::A").unwrap(), SymbolKind::Struct)
            .unwrap();
        let b_id = registry
            .register(SymbolPath::parse("test::B").unwrap(), SymbolKind::Struct)
            .unwrap();

        let mut graph = TypeFlowGraphV2::new();
        graph.add_definition(a_id, TypeDefKind::Struct);
        graph.add_definition(b_id, TypeDefKind::Struct);

        // A uses B, B uses A (cycle)
        graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            Some(b_id),
            Some(a_id),
        );
        graph.add_usage(
            UsageContext::FieldType,
            RefKind::Owned,
            Some(a_id),
            Some(b_id),
        );

        // Should not infinite loop; visited set prevents revisiting
        let chain = graph.type_users_chain(a_id, 10);
        assert_eq!(chain.len(), 1); // Only B at depth 1
        assert_eq!(chain[0].symbol, b_id);
    }
}