pmat 3.29.0

PMAT - Zero-config AI context generation and code quality toolkit (CLI, MCP)
impl DependencyGraph {
    #[must_use]
    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    /// Create a new instance.
    pub fn new() -> Self {
        Self {
            nodes: FxHashMap::default(),
            edges: Vec::new(),
        }
    }

    /// Adds a node to the dependency graph
    ///
    /// # Examples
    ///
    /// ```rust
    /// use pmat::models::dag::{DependencyGraph, NodeInfo, NodeType};
    /// use rustc_hash::FxHashMap;
    ///
    /// let mut graph = DependencyGraph::new();
    /// graph.add_node(NodeInfo {
    ///     id: "main::hello".to_string(),
    ///     label: "hello".to_string(),
    ///     node_type: NodeType::Function,
    ///     file_path: "src/main.rs".to_string(),
    ///     line_number: 10,
    ///     complexity: 1,
    ///     metadata: FxHashMap::default(),
    /// });
    ///
    /// assert_eq!(graph.nodes.len(), 1);
    /// ```
    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    pub fn add_node(&mut self, node: NodeInfo) {
        self.nodes.insert(node.id.clone(), node);
    }

    /// Get the number of nodes in the graph
    #[must_use]
    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }

    /// Get the number of edges in the graph  
    #[must_use]
    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    pub fn edge_count(&self) -> usize {
        self.edges.len()
    }

    /// Adds an edge between two nodes in the dependency graph
    ///
    /// # Examples
    ///
    /// ```rust
    /// use pmat::models::dag::{DependencyGraph, Edge, EdgeType};
    ///
    /// let mut graph = DependencyGraph::new();
    /// graph.add_edge(Edge {
    ///     from: "main::hello".to_string(),
    ///     to: "utils::print".to_string(),
    ///     edge_type: EdgeType::Calls,
    ///     weight: 1,
    /// });
    ///
    /// assert_eq!(graph.edges.len(), 1);
    /// assert_eq!(graph.edges[0].edge_type, EdgeType::Calls);
    /// ```
    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    pub fn add_edge(&mut self, edge: Edge) {
        self.edges.push(edge);
    }

    /// Creates a new graph containing only edges of the specified type
    ///
    /// # Examples
    ///
    /// ```rust
    /// use pmat::models::dag::{DependencyGraph, Edge, EdgeType};
    ///
    /// let mut graph = DependencyGraph::new();
    /// graph.add_edge(Edge {
    ///     from: "a".to_string(),
    ///     to: "b".to_string(),
    ///     edge_type: EdgeType::Calls,
    ///     weight: 1,
    /// });
    /// graph.add_edge(Edge {
    ///     from: "c".to_string(),
    ///     to: "d".to_string(),
    ///     edge_type: EdgeType::Imports,
    ///     weight: 1,
    /// });
    ///
    /// let calls_only = graph.filter_by_edge_type(EdgeType::Calls);
    /// assert_eq!(calls_only.edges.len(), 1);
    /// assert_eq!(calls_only.edges[0].edge_type, EdgeType::Calls);
    /// ```
    #[must_use]
    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    pub fn filter_by_edge_type(&self, edge_type: EdgeType) -> Self {
        let filtered_edges: Vec<Edge> = self
            .edges
            .iter()
            .filter(|e| e.edge_type == edge_type)
            .cloned()
            .collect();

        // If filtering results in no edges but we originally had edges,
        // only include nodes that were connected by the filtered edge type
        if filtered_edges.is_empty() && !self.edges.is_empty() {
            // Return empty nodes since no nodes are connected by this edge type
            return Self {
                nodes: FxHashMap::default(),
                edges: filtered_edges,
            };
        }

        // If we have no edges at all, return all nodes
        if self.edges.is_empty() {
            return Self {
                nodes: self.nodes.clone(),
                edges: Vec::new(),
            };
        }

        // Otherwise, filter nodes to only those connected by the filtered edges
        let used_nodes: FxHashSet<String> = filtered_edges
            .iter()
            .flat_map(|e| vec![e.from.clone(), e.to.clone()])
            .collect();

        let filtered_nodes: FxHashMap<String, NodeInfo> = self
            .nodes
            .iter()
            .filter(|(id, _)| used_nodes.contains(*id))
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();

        Self {
            nodes: filtered_nodes,
            edges: filtered_edges,
        }
    }

    /// Creates the sub-graph induced by a set of edge types
    ///
    /// Unlike [`Self::filter_by_edge_type`], the result never falls back to
    /// "all nodes": a graph asked for its call edges either contains call edges and
    /// the functions they connect, or is empty. `analyze --dag-type` relies on this
    /// to make its four graph types genuinely different (#653).
    ///
    /// # Examples
    ///
    /// ```rust
    /// use pmat::models::dag::{DependencyGraph, Edge, EdgeType, NodeInfo, NodeType};
    /// use rustc_hash::FxHashMap;
    ///
    /// let mut graph = DependencyGraph::new();
    /// for id in ["a", "b", "c"] {
    ///     graph.add_node(NodeInfo {
    ///         id: id.to_string(),
    ///         label: id.to_string(),
    ///         node_type: NodeType::Function,
    ///         file_path: "src/lib.rs".to_string(),
    ///         line_number: 1,
    ///         complexity: 1,
    ///         metadata: FxHashMap::default(),
    ///     });
    /// }
    /// graph.add_edge(Edge { from: "a".into(), to: "b".into(), edge_type: EdgeType::Calls, weight: 1 });
    /// graph.add_edge(Edge { from: "b".into(), to: "c".into(), edge_type: EdgeType::Imports, weight: 1 });
    ///
    /// let calls = graph.filter_by_edge_types(&[EdgeType::Calls]);
    /// assert_eq!(calls.edges.len(), 1);
    /// assert_eq!(calls.nodes.len(), 2); // only a and b
    /// ```
    #[must_use]
    pub fn filter_by_edge_types(&self, edge_types: &[EdgeType]) -> Self {
        let filtered_edges: Vec<Edge> = self
            .edges
            .iter()
            .filter(|e| edge_types.contains(&e.edge_type))
            .cloned()
            .collect();

        let used_nodes: FxHashSet<&String> = filtered_edges
            .iter()
            .flat_map(|e| [&e.from, &e.to])
            .collect();

        let filtered_nodes: FxHashMap<String, NodeInfo> = self
            .nodes
            .iter()
            .filter(|(id, _)| used_nodes.contains(id))
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();

        Self {
            nodes: filtered_nodes,
            edges: filtered_edges,
        }
    }
}

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