graphrust_core/

graph.rs

//! Core graph data structure.
//!
//! Provides an efficient adjacency list representation optimized for:
//! - Cache locality (dense packing of edges)
//! - FFI compatibility (repr(C) where beneficial)
//! - Sparse graph optimization (only stores edges that exist)
//! - Both directed and undirected graphs

use crate::types::{EdgeWeight, NodeId, NodeWeight};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// A directed or undirected graph using adjacency list representation.
///
/// # Storage Layout
/// - Nodes stored in vec for O(1) access and cache locality
/// - Edges stored as Vec<Vec<NodeId>> for adjacency list
/// - Optional weights stored in separate hash maps to minimize memory when unweighted
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Graph {
    /// Number of nodes in the graph
    num_nodes: u32,
    /// Number of edges (for directed graphs, this counts each directed edge once)
    num_edges: u32,
    /// Adjacency list: for each node, list of neighbors (outgoing edges in directed graphs)
    adjacency: Vec<Vec<NodeId>>,
    /// Optional node weights (e.g., for PageRank initial values)
    node_weights: Option<Vec<NodeWeight>>,
    /// Optional edge weights (e.g., for Dijkstra shortest paths)
    edge_weights: Option<HashMap<(NodeId, NodeId), EdgeWeight>>,
    /// Whether the graph is directed
    is_directed: bool,
}

impl Graph {
    /// Creates a new empty graph with capacity for `num_nodes` nodes.
    pub fn new(num_nodes: u32, is_directed: bool) -> Self {
        Graph {
            num_nodes,
            num_edges: 0,
            adjacency: vec![Vec::new(); num_nodes as usize],
            node_weights: None,
            edge_weights: None,
            is_directed,
        }
    }

    /// Returns the number of nodes in the graph.
    #[inline]
    pub fn num_nodes(&self) -> u32 {
        self.num_nodes
    }

    /// Returns the number of edges in the graph.
    #[inline]
    pub fn num_edges(&self) -> u32 {
        self.num_edges
    }

    /// Returns whether the graph is directed.
    #[inline]
    pub fn is_directed(&self) -> bool {
        self.is_directed
    }

    /// Adds an edge from `from` to `to` with optional weight.
    ///
    /// For undirected graphs, this automatically adds the reverse edge.
    ///
    /// # Panics
    /// Panics if `from` or `to` are >= `num_nodes`.
    pub fn add_edge(&mut self, from: NodeId, to: NodeId, weight: Option<EdgeWeight>) {
        assert!(from.0 < self.num_nodes, "Node {} out of bounds", from.0);
        assert!(to.0 < self.num_nodes, "Node {} out of bounds", to.0);

        // Check if edge already exists
        if !self.adjacency[from.as_usize()].contains(&to) {
            self.adjacency[from.as_usize()].push(to);
            self.num_edges += 1;
        }

        if let Some(w) = weight {
            if self.edge_weights.is_none() {
                self.edge_weights = Some(HashMap::new());
            }
            self.edge_weights
                .as_mut()
                .unwrap()
                .insert((from, to), w);
        }

        // For undirected graphs, add reverse edge
        if !self.is_directed && from != to {
            if !self.adjacency[to.as_usize()].contains(&from) {
                self.adjacency[to.as_usize()].push(from);
                // Note: num_edges already incremented once; undirected graphs count each edge once
            }
            if let Some(w) = weight {
                self.edge_weights
                    .as_mut()
                    .unwrap()
                    .insert((to, from), w);
            }
        }
    }

    /// Gets the neighbors (outgoing edges) for a node.
    #[inline]
    pub fn neighbors(&self, node: NodeId) -> &[NodeId] {
        &self.adjacency[node.as_usize()]
    }

    /// Gets the weight of an edge.
    pub fn edge_weight(&self, from: NodeId, to: NodeId) -> Option<EdgeWeight> {
        self.edge_weights
            .as_ref()
            .and_then(|weights| weights.get(&(from, to)).copied())
            .or(Some(EdgeWeight::default()))
    }

    /// Sets weights for all nodes.
    pub fn set_node_weights(&mut self, weights: Vec<NodeWeight>) {
        assert_eq!(
            weights.len() as u32, self.num_nodes,
            "Node weights length must match num_nodes"
        );
        self.node_weights = Some(weights);
    }

    /// Gets the weight of a node.
    pub fn node_weight(&self, node: NodeId) -> Option<NodeWeight> {
        self.node_weights
            .as_ref()
            .map(|weights| weights[node.as_usize()])
    }

    /// Iterates over all nodes in the graph.
    pub fn nodes(&self) -> impl Iterator<Item = NodeId> {
        (0..self.num_nodes).map(NodeId::new)
    }

    /// Iterates over all edges in the graph as (from, to) tuples.
    pub fn edges(&self) -> impl Iterator<Item = (NodeId, NodeId)> + '_ {
        self.adjacency
            .iter()
            .enumerate()
            .flat_map(|(from_idx, neighbors)| {
                let from = NodeId(from_idx as u32);
                neighbors.iter().map(move |&to| (from, to))
            })
    }
}

/// Builder pattern for constructing graphs.
#[derive(Default)]
pub struct GraphBuilder {
    num_nodes: u32,
    is_directed: bool,
    edges: Vec<(NodeId, NodeId, Option<EdgeWeight>)>,
}

impl GraphBuilder {
    /// Creates a new graph builder.
    pub fn new(num_nodes: u32) -> Self {
        GraphBuilder {
            num_nodes,
            is_directed: false,
            edges: Vec::new(),
        }
    }

    /// Sets whether the graph is directed.
    pub fn directed(mut self, directed: bool) -> Self {
        self.is_directed = directed;
        self
    }

    /// Adds an edge to the builder.
    pub fn add_edge(mut self, from: NodeId, to: NodeId) -> Self {
        self.edges.push((from, to, None));
        self
    }

    /// Adds a weighted edge to the builder.
    pub fn add_weighted_edge(mut self, from: NodeId, to: NodeId, weight: f64) -> Self {
        self.edges.push((from, to, Some(EdgeWeight::new(weight))));
        self
    }

    /// Builds the graph.
    pub fn build(self) -> Graph {
        let mut graph = Graph::new(self.num_nodes, self.is_directed);
        for (from, to, weight) in self.edges {
            graph.add_edge(from, to, weight);
        }
        graph
    }
}

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

    #[test]
    fn test_graph_creation() {
        let graph = Graph::new(5, true);
        assert_eq!(graph.num_nodes(), 5);
        assert_eq!(graph.num_edges(), 0);
        assert!(graph.is_directed());
    }

    #[test]
    fn test_add_edge() {
        let mut graph = Graph::new(5, true);
        graph.add_edge(NodeId(0), NodeId(1), None);
        assert_eq!(graph.num_edges(), 1);
        assert!(graph.neighbors(NodeId(0)).contains(&NodeId(1)));
    }

    #[test]
    fn test_undirected_graph() {
        let mut graph = Graph::new(5, false);
        graph.add_edge(NodeId(0), NodeId(1), None);
        assert_eq!(graph.num_edges(), 1);
        assert!(graph.neighbors(NodeId(0)).contains(&NodeId(1)));
        assert!(graph.neighbors(NodeId(1)).contains(&NodeId(0)));
    }

    #[test]
    fn test_builder_pattern() {
        let graph = GraphBuilder::new(4)
            .directed(true)
            .add_edge(NodeId(0), NodeId(1))
            .add_weighted_edge(NodeId(1), NodeId(2), 2.5)
            .build();

        assert_eq!(graph.num_nodes(), 4);
        assert_eq!(graph.num_edges(), 2);
        assert_eq!(graph.edge_weight(NodeId(1), NodeId(2)), Some(EdgeWeight(2.5)));
    }

    #[test]
    fn test_node_weights() {
        let mut graph = Graph::new(3, false);
        graph.set_node_weights(vec![
            NodeWeight(1.0),
            NodeWeight(2.0),
            NodeWeight(3.0),
        ]);

        assert_eq!(graph.node_weight(NodeId(0)), Some(NodeWeight(1.0)));
        assert_eq!(graph.node_weight(NodeId(1)), Some(NodeWeight(2.0)));
    }
}