spokes 0.3.0

//! Algorithms for various network related tasks.

use std::{
    collections::{HashSet, VecDeque},
    fmt::Debug,
    hash::Hash,
    marker::PhantomData,
};

use crate::{ArcInfo, arc_storage::ArcStorage};

/// Direction for seaching with a searching algorithm
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Direction {
    /// Only include forward arcs
    Forward,
    /// Only include reverse arcs
    Reverse,
    /// Include both forward and reverse arcs
    Undirected,
}

/// Iterator for Depth first traversal of a network starting at node `starting_node`.
///
/// # Example
/// ```rust
/// use spokes::search::{DfsIter, Direction};
/// use spokes::arc_storage::{ForwardAndReverseStar, ArcStorage};
///
/// // A B C D E F G
/// // 0 1 2 3 4 5 6
/// let (a, b, c, d, e, f, g) = (0, 1, 2, 3, 4, 5, 6);
/// let mut arcs = ForwardAndReverseStar::new();
/// arcs.add_arcs([
///     (a, b), (b, a),
///     (a, c), (c, a),
///     (a, e), (e, a),
///     (b, d), (d, b),
///     (b, f), (f, b),
///     (c, g), (g, c),
///     (e, f), (f, e),
/// ]);
///
/// assert_eq!(
///     DfsIter::new(&arcs, &0, Direction::Forward).copied().collect::<Vec<usize>>(),
///     vec![a, e, f, b, d, c, g],
/// );
/// ```
pub struct DfsIter<'a, I, A, ArcStore>
where
    ArcStore: ArcStorage<I, A>,
{
    arcs: &'a ArcStore,
    to_visit: VecDeque<&'a I>,
    labeled: HashSet<&'a I>,
    direction: Direction,
    _phantom_a: PhantomData<A>,
}

impl<'a, I, A, AS> DfsIter<'a, I, A, AS>
where
    AS: ArcStorage<I, A>,
{
    /// Create a new Depth-First Iterator of the arcs in `arcs` starting from `starting_node`.
    #[must_use]
    pub fn new(arcs: &'a AS, starting_node: &'a I, direction: Direction) -> Self {
        let mut to_visit = VecDeque::new();
        to_visit.push_front(starting_node);
        Self {
            arcs,
            to_visit,
            direction,
            labeled: HashSet::new(),
            _phantom_a: PhantomData,
        }
    }
}

impl<'a, I, A, AS> Iterator for DfsIter<'a, I, A, AS>
where
    AS: ArcStorage<I, A>,
    I: Hash + Eq + Copy,
    A: 'a,
{
    type Item = &'a I;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(v) = self.to_visit.pop_front() {
            if !self.labeled.contains(&v) {
                self.labeled.insert(v);

                let edges =
                    match self.direction {
                        Direction::Forward => &mut self.arcs.forward_arcs(v)
                            as &mut dyn Iterator<Item = &ArcInfo<I, A>>,
                        Direction::Reverse => &mut self.arcs.reverse_arcs(v)
                            as &mut dyn Iterator<Item = &ArcInfo<I, A>>,
                        Direction::Undirected => {
                            &mut self.arcs.all_arcs(v) as &mut dyn Iterator<Item = &ArcInfo<I, A>>
                        }
                    };

                for edge in edges {
                    match self.direction {
                        Direction::Forward => {
                            if !self.labeled.contains(&edge.head) {
                                self.to_visit.push_front(&edge.head);
                            }
                        }
                        Direction::Reverse => {
                            if !self.labeled.contains(&edge.tail) {
                                self.to_visit.push_front(&edge.tail);
                            }
                        }
                        Direction::Undirected => {
                            if !self.labeled.contains(&edge.head) {
                                self.to_visit.push_front(&edge.head);
                            }
                            if !self.labeled.contains(&edge.tail) {
                                self.to_visit.push_front(&edge.tail);
                            }
                        }
                    }
                }
                return Some(v);
            }
        }

        None
    }
}

/// Iterator for Breadth first traversal of a network starting at node `starting_node`.
///
/// # Example
/// ```rust
/// use spokes::search::{BfsIter, Direction};
/// use spokes::arc_storage::{ForwardAndReverseStar, ArcStorage};
///
/// let (a, b, c, d, e, f, g) = (0, 1, 2, 3, 4, 5, 6);
/// let mut arcs = ForwardAndReverseStar::new();
/// arcs.add_arcs([
///     (a, b), (b, a),
///     (a, c), (c, a),
///     (a, e), (e, a),
///     (b, d), (d, b),
///     (b, f), (f, b),
///     (c, g), (g, c),
///     (e, f), (f, e),
/// ]);
///
/// assert_eq!(
///     BfsIter::new(&arcs, &a, Direction::Forward).copied().collect::<Vec<usize>>(),
///     vec![a, b, c, e, d, f, g],
/// );
/// ```
pub struct BfsIter<'a, I, A, ArcStore>
where
    ArcStore: ArcStorage<I, A>,
{
    arcs: &'a ArcStore,
    to_visit: VecDeque<&'a I>,
    labeled: HashSet<&'a I>,
    direction: Direction,
    _phantom_a: PhantomData<A>,
}

impl<'a, I, A, AS> BfsIter<'a, I, A, AS>
where
    AS: ArcStorage<I, A>,
{
    /// Create a new Breadth-First Iterator of the arcs in `arcs` starting from `starting_node`.
    #[must_use]
    pub fn new(arcs: &'a AS, starting_node: &'a I, direction: Direction) -> Self {
        let mut to_visit = VecDeque::new();
        to_visit.push_front(starting_node);
        Self {
            arcs,
            to_visit,
            labeled: HashSet::new(),
            direction,
            _phantom_a: PhantomData,
        }
    }
}

impl<'a, I, A, AS> Iterator for BfsIter<'a, I, A, AS>
where
    AS: ArcStorage<I, A>,
    I: Hash + Eq + Copy,
    A: 'a,
{
    type Item = &'a I;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(v) = self.to_visit.pop_front() {
            if !self.labeled.contains(&v) {
                self.labeled.insert(v);

                let edges =
                    match self.direction {
                        Direction::Forward => &mut self.arcs.forward_arcs(v)
                            as &mut dyn Iterator<Item = &ArcInfo<I, A>>,
                        Direction::Reverse => &mut self.arcs.reverse_arcs(v)
                            as &mut dyn Iterator<Item = &ArcInfo<I, A>>,
                        Direction::Undirected => {
                            &mut self.arcs.all_arcs(v) as &mut dyn Iterator<Item = &ArcInfo<I, A>>
                        }
                    };

                for edge in edges {
                    match self.direction {
                        Direction::Forward => {
                            if !self.labeled.contains(&edge.head) {
                                self.to_visit.push_back(&edge.head);
                            }
                        }
                        Direction::Reverse => {
                            if !self.labeled.contains(&edge.tail) {
                                self.to_visit.push_back(&edge.tail);
                            }
                        }
                        Direction::Undirected => {
                            if !self.labeled.contains(&edge.head) {
                                self.to_visit.push_back(&edge.head);
                            }
                            if !self.labeled.contains(&edge.tail) {
                                self.to_visit.push_back(&edge.tail);
                            }
                        }
                    }
                }

                return Some(v);
            }
        }

        None
    }
}

#[cfg(test)]
mod tests {
    use crate::Network;

    use super::*;

    #[test]
    fn simple_bfs() {
        let mut network: Network<i32, (), ()> = Network::new();
        network.add_nodes((0..4).map(|i| (i, ())));

        network
            .add_arcs([(0, 1, ()), (1, 2, ()), (2, 3, ())])
            .unwrap();
        assert_eq!(
            BfsIter::new(network.arc_store(), &0, Direction::Forward)
                .copied()
                .collect::<Vec<i32>>(),
            vec![0, 1, 2, 3]
        );

        assert_eq!(
            BfsIter::new(network.arc_store(), &3, Direction::Reverse)
                .copied()
                .collect::<Vec<i32>>(),
            vec![3, 2, 1, 0]
        );

        assert_eq!(
            BfsIter::new(network.arc_store(), &2, Direction::Undirected)
                .copied()
                .collect::<Vec<i32>>(),
            vec![2, 3, 1, 0]
        );
    }
}