rustsim_spaces/

graph.rs

//! Graph-based discrete space.
//!
//! [`GraphSpace`] represents a network where each node can hold agents.
//! Agents move along edges, and distance is measured in hops (BFS).
//! Supports both directed and undirected graphs with dynamic topology
//! (vertices and edges can be added/removed at runtime).
//!
//! Mirrors Julia Agents.jl `GraphSpace`.

use rand::Rng;
use rustsim_core::{
    interaction::{PositionedAgent, SpaceInteraction},
    space::Space,
    types::{AgentId, NodeId},
};
use std::collections::{HashSet, VecDeque};
use thiserror::Error;

/// Position in a graph space - a node index (`0..num_vertices`).
pub type GraphPos = NodeId;

/// Errors returned by graph space operations.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum GraphSpaceError {
    /// The node index is out of range.
    #[error("invalid graph node index {0}")]
    InvalidNode(GraphPos),
}

/// Neighbor search direction for directed graphs.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum NeighborType {
    /// Outgoing neighbors (default; equivalent to All for undirected graphs).
    #[default]
    Out,
    /// Incoming neighbors.
    In,
    /// Both incoming and outgoing.
    All,
}

/// A graph-based discrete space mirroring Julia Agents.jl `GraphSpace`.
///
/// Each node can hold an arbitrary number of agents. Agents move between
/// nodes along edges. Distance is measured in graph hops.
///
/// Supports both directed and undirected graphs. Edges and vertices can
/// be added/removed dynamically at runtime.
///
/// # Example: building floor graph
/// ```text
///   [Floor1-RoomA] --edge-- [Floor1-Corridor] --edge-- [Floor1-RoomB]
///                                |
///                              (stairs)
///                                |
///   [Floor2-RoomA] --edge-- [Floor2-Corridor] --edge-- [Floor2-RoomB]
/// ```
/// Each room/corridor/stairwell is a node. Edges connect adjacent spaces.
#[derive(Debug, Clone)]
pub struct GraphSpace {
    /// Adjacency list: out-neighbors for each node.
    adj_out: Vec<Vec<GraphPos>>,
    /// Adjacency list: in-neighbors for each node (for directed graphs).
    adj_in: Vec<Vec<GraphPos>>,
    /// Agent IDs stored at each node.
    stored_ids: Vec<Vec<AgentId>>,
    /// Whether the graph is directed.
    directed: bool,
}

impl GraphSpace {
    /// Create an undirected graph with `n` nodes and no edges.
    pub fn new(n: usize) -> Self {
        Self {
            adj_out: vec![Vec::new(); n],
            adj_in: vec![Vec::new(); n],
            stored_ids: vec![Vec::new(); n],
            directed: false,
        }
    }

    /// Create a graph with `n` nodes. If `directed` is true, edges are one-way.
    pub fn new_directed(n: usize, directed: bool) -> Self {
        Self {
            adj_out: vec![Vec::new(); n],
            adj_in: vec![Vec::new(); n],
            stored_ids: vec![Vec::new(); n],
            directed,
        }
    }

    /// Number of nodes in the graph.
    pub fn num_vertices(&self) -> usize {
        self.adj_out.len()
    }

    /// Number of edges in the graph.
    ///
    /// For undirected graphs, each edge is counted once (not twice).
    pub fn num_edges(&self) -> usize {
        let total: usize = self.adj_out.iter().map(|v| v.len()).sum();
        if self.directed {
            total
        } else {
            total / 2
        }
    }

    /// Whether the graph is directed.
    pub fn is_directed(&self) -> bool {
        self.directed
    }

    /// Add a new vertex and return its index.
    pub fn add_vertex(&mut self) -> GraphPos {
        let idx = self.adj_out.len();
        self.adj_out.push(Vec::new());
        self.adj_in.push(Vec::new());
        self.stored_ids.push(Vec::new());
        idx
    }

    /// Remove a vertex by swapping with the last vertex (matches Graphs.jl behavior).
    /// All agents at the removed vertex must be removed beforehand.
    /// Returns true if successful.
    pub fn rem_vertex(&mut self, n: GraphPos) -> bool {
        let nv = self.num_vertices();
        if n >= nv {
            return false;
        }

        // Remove all edges involving node n
        let out_neighbors: Vec<GraphPos> = self.adj_out[n].clone();
        for &neighbor in &out_neighbors {
            self.rem_edge(n, neighbor);
        }
        let in_neighbors: Vec<GraphPos> = self.adj_in[n].clone();
        for &neighbor in &in_neighbors {
            self.rem_edge(neighbor, n);
        }

        let last = nv - 1;
        if n != last {
            // Swap last node into position n
            self.adj_out.swap(n, last);
            self.adj_in.swap(n, last);
            self.stored_ids.swap(n, last);

            // Update all references from `last` to `n`
            for neighbors in &mut self.adj_out {
                for pos in neighbors.iter_mut() {
                    if *pos == last {
                        *pos = n;
                    }
                }
            }
            for neighbors in &mut self.adj_in {
                for pos in neighbors.iter_mut() {
                    if *pos == last {
                        *pos = n;
                    }
                }
            }
        }

        self.adj_out.pop();
        self.adj_in.pop();
        self.stored_ids.pop();
        true
    }

    /// Add an edge from `a` to `b`. For undirected graphs, also adds `b` to `a`.
    pub fn add_edge(&mut self, a: GraphPos, b: GraphPos) -> bool {
        let nv = self.num_vertices();
        if a >= nv || b >= nv {
            return false;
        }
        if self.adj_out[a].contains(&b) {
            return false; // already exists
        }
        self.adj_out[a].push(b);
        self.adj_in[b].push(a);
        if !self.directed {
            self.adj_out[b].push(a);
            self.adj_in[a].push(b);
        }
        true
    }

    /// Remove an edge from `a` to `b`.
    pub fn rem_edge(&mut self, a: GraphPos, b: GraphPos) -> bool {
        let nv = self.num_vertices();
        if a >= nv || b >= nv {
            return false;
        }
        let removed = remove_from_vec(&mut self.adj_out[a], b);
        remove_from_vec(&mut self.adj_in[b], a);
        if !self.directed {
            remove_from_vec(&mut self.adj_out[b], a);
            remove_from_vec(&mut self.adj_in[a], b);
        }
        removed
    }

    /// Outgoing neighbors of a node (slice reference, no allocation).
    pub fn neighbors_out(&self, n: GraphPos) -> &[GraphPos] {
        &self.adj_out[n]
    }

    /// Incoming neighbors of a node (slice reference, no allocation).
    pub fn neighbors_in(&self, n: GraphPos) -> &[GraphPos] {
        &self.adj_in[n]
    }

    /// All neighbors (union of incoming and outgoing), deduplicated.
    pub fn neighbors_all(&self, n: GraphPos) -> Vec<GraphPos> {
        let mut set: HashSet<GraphPos> = HashSet::new();
        set.extend(&self.adj_out[n]);
        set.extend(&self.adj_in[n]);
        set.into_iter().collect()
    }

    /// Get neighbors according to a [`NeighborType`] selector.
    pub fn neighbors(&self, n: GraphPos, kind: NeighborType) -> Vec<GraphPos> {
        match kind {
            NeighborType::Out => self.adj_out[n].clone(),
            NeighborType::In => self.adj_in[n].clone(),
            NeighborType::All => self.neighbors_all(n),
        }
    }

    /// Agent IDs stored at a given node.
    pub fn ids_in_position(&self, n: GraphPos) -> &[AgentId] {
        &self.stored_ids[n]
    }

    /// All valid node indices.
    pub fn positions(&self) -> std::ops::Range<usize> {
        0..self.num_vertices()
    }

    /// Find all node indices reachable within `r` hops via BFS, excluding the origin.
    pub fn nearby_positions(&self, pos: GraphPos, r: usize, kind: NeighborType) -> Vec<GraphPos> {
        let mut visited = HashSet::new();
        let mut queue = VecDeque::new();
        visited.insert(pos);
        queue.push_back((pos, 0usize));

        let mut result = Vec::new();

        while let Some((node, dist)) = queue.pop_front() {
            if dist > 0 {
                result.push(node);
            }
            if dist < r {
                match kind {
                    NeighborType::Out => {
                        for &neighbor in &self.adj_out[node] {
                            if visited.insert(neighbor) {
                                queue.push_back((neighbor, dist + 1));
                            }
                        }
                    }
                    NeighborType::In => {
                        for &neighbor in &self.adj_in[node] {
                            if visited.insert(neighbor) {
                                queue.push_back((neighbor, dist + 1));
                            }
                        }
                    }
                    NeighborType::All => {
                        for &neighbor in &self.adj_out[node] {
                            if visited.insert(neighbor) {
                                queue.push_back((neighbor, dist + 1));
                            }
                        }
                        for &neighbor in &self.adj_in[node] {
                            if visited.insert(neighbor) {
                                queue.push_back((neighbor, dist + 1));
                            }
                        }
                    }
                }
            }
        }
        result
    }

    /// Find all agent IDs within `r` hops, including agents at the origin node.
    pub fn nearby_agent_ids(&self, pos: GraphPos, r: usize, kind: NeighborType) -> Vec<AgentId> {
        let mut ids = Vec::new();
        // Include agents at origin
        ids.extend_from_slice(&self.stored_ids[pos]);
        // Include agents at nearby positions
        for neighbor in self.nearby_positions(pos, r, kind) {
            ids.extend_from_slice(&self.stored_ids[neighbor]);
        }
        ids
    }
}

fn remove_from_vec(v: &mut Vec<GraphPos>, val: GraphPos) -> bool {
    if let Some(i) = v.iter().position(|&x| x == val) {
        v.swap_remove(i);
        true
    } else {
        false
    }
}

impl Space for GraphSpace {}

impl<A> SpaceInteraction<A> for GraphSpace
where
    A: PositionedAgent<Position = GraphPos>,
{
    type Error = GraphSpaceError;

    fn random_position<R: rand::RngCore>(&self, rng: &mut R) -> A::Position {
        rng.gen_range(0..self.num_vertices())
    }

    fn add_agent(&mut self, agent: &A) -> Result<(), Self::Error> {
        let pos = *agent.position();
        if pos >= self.num_vertices() {
            return Err(GraphSpaceError::InvalidNode(pos));
        }
        self.stored_ids[pos].push(agent.id());
        Ok(())
    }

    fn remove_agent(&mut self, agent: &A) -> Result<(), Self::Error> {
        let pos = *agent.position();
        if pos >= self.num_vertices() {
            return Err(GraphSpaceError::InvalidNode(pos));
        }
        if let Some(i) = self.stored_ids[pos].iter().position(|&id| id == agent.id()) {
            self.stored_ids[pos].swap_remove(i);
        }
        Ok(())
    }

    fn nearby_ids(&self, position: &A::Position, radius: usize) -> Vec<AgentId> {
        self.nearby_agent_ids(*position, radius, NeighborType::Out)
    }
}