spokes 0.4.0-rc.4

//! Algorithms for various network related tasks.

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

use crate::{ArcInfo, Network, arc_storage::ArcStorage, node_storage::NodeStorage};

/// 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, &a, 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,
    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.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,
    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.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
    }
}

/// Create an iterator over all all non-cycle paths in the network.
pub struct PathIter<'a, I, A, ArcStore>
where
    ArcStore: ArcStorage<I, A>,
{
    arcs: &'a ArcStore,
    curent_path: Vec<&'a I>,
    branches: VecDeque<(usize, &'a I)>,
    seen: HashSet<&'a I>,
    direction: Direction,
    _phantom_a: PhantomData<A>,
}

impl<'a, I, A, AS> PathIter<'a, I, A, AS>
where
    A: 'a,
    AS: ArcStorage<I, A>,
    I: Eq + Hash,
{
    /// 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 branches = match direction {
            Direction::Forward => arcs
                .forward_arcs(starting_node)
                .map(|arc| (1, arc.head()))
                .collect(),
            Direction::Reverse => arcs
                .reverse_arcs(starting_node)
                .map(|arc| (1, arc.tail()))
                .collect(),
            Direction::Undirected => arcs
                .reverse_arcs(starting_node)
                .map(|arc| (1, arc.tail()))
                .chain(arcs.forward_arcs(starting_node).map(|arc| (1, arc.head())))
                .collect(),
        };

        Self {
            arcs,
            curent_path: vec![starting_node],
            branches,
            direction,
            seen: HashSet::from([starting_node]),
            _phantom_a: PhantomData,
        }
    }
}

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

    fn next(&mut self) -> Option<Self::Item> {
        if let Some((len, next_node)) = self.branches.pop_front() {
            self.curent_path.drain(len..).for_each(|x| {
                self.seen.remove(x);
            });
            self.curent_path.push(next_node);
            self.seen.insert(next_node);

            match self.direction {
                Direction::Forward => self.arcs.forward_arcs(next_node).for_each(|arc| {
                    if !self.seen.contains(arc.head()) {
                        self.branches.push_front((len + 1, arc.head()));
                    }
                }),
                Direction::Reverse => self.arcs.reverse_arcs(next_node).for_each(|arc| {
                    if !self.seen.contains(arc.tail()) {
                        self.branches.push_front((len + 1, arc.tail()));
                    }
                }),
                Direction::Undirected => {
                    self.arcs.forward_arcs(next_node).for_each(|arc| {
                        if !self.seen.contains(arc.head()) {
                            self.branches.push_front((len + 1, arc.head()));
                        }
                    });
                    self.arcs.reverse_arcs(next_node).for_each(|arc| {
                        if !self.seen.contains(arc.tail()) {
                            self.branches.push_front((len + 1, arc.tail()));
                        }
                    });
                }
            }

            Some(self.curent_path.clone())
        } else {
            None
        }
    }
}

pub(crate) fn intermediate_nodes<'a, 'b, I, NodeAttr, ArcAttr, NodeStore, ArcStore>(
    network: &'a Network<I, NodeAttr, ArcAttr, NodeStore, ArcStore>,
    ends: impl IntoIterator<Item = &'b I>,
    direction: Direction,
) -> HashSet<&'a I>
where
    I: Hash + Eq + 'b,
    NodeStore: NodeStorage<I, NodeAttr>,
    ArcStore: ArcStorage<I, ArcAttr>,
    'b: 'a,
{
    let ends: HashSet<_> = ends.into_iter().collect();
    let mut output: HashSet<_> = ends.clone();

    for &start_id in &ends {
        for path in PathIter::new(network.arc_store(), start_id, direction) {
            let last_added = *path
                .last()
                .expect("At least one element should be in a path");

            if ends.contains(last_added) {
                for id in path {
                    output.insert(id);
                }
            }
        }
    }

    output
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used, clippy::zero_sized_map_values)]

    use super::*;
    use crate::{HashMapStorage, Network};
    use std::collections::HashMap;

    #[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, ())])
            .expect("should construct");

        println!("{}", network.dot().to_string().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]
        );
    }

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

        network
            .add_arcs([
                (0, 1),
                (0, 2),
                (0, 3),
                (1, 4),
                (2, 4),
                (3, 4),
                (0, 5),
                (4, 6),
                (1, 7),
            ])
            .expect("should construct");

        let mut paths: Vec<Vec<u8>> = PathIter::new(network.arc_store(), &0, Direction::Forward)
            .map(|x| x.into_iter().copied().collect())
            .collect();
        paths.sort_by_key(|x| {
            x.iter()
                .enumerate()
                .map(|(i, v)| (*v as usize) << i)
                .sum::<usize>()
        });

        let mut expected = vec![
            vec![0_u8, 1],
            vec![0, 1, 7],
            vec![0, 1, 4],
            vec![0, 1, 4, 6],
            vec![0, 5],
            vec![0, 3],
            vec![0, 3, 4],
            vec![0, 3, 4, 6],
            vec![0, 2],
            vec![0, 2, 4],
            vec![0, 2, 4, 6],
        ];
        expected.sort_by_key(|x| {
            x.iter()
                .enumerate()
                .map(|(i, v)| (*v as usize) << i)
                .sum::<usize>()
        });

        assert_eq!(paths, expected);
    }

    #[test]
    fn simple_intermediate_nodes() {
        let mut net: Network<char, (), (), HashMap<char, ()>, HashMapStorage<char, ()>> =
            Network::new();

        net.add_nodes([('A', ()), ('B', ()), ('C', ())]);
        net.add_arcs([('A', 'B', ()), ('A', 'C', ())]).unwrap();

        let in_nodes = intermediate_nodes(&net, [&'B', &'C'], Direction::Undirected);

        assert_eq!(in_nodes, HashSet::from([&'A', &'B', &'C']));
    }

    #[test]
    fn elaborate_intermediate_nodes() {
        let mut net: Network<char, (), (), HashMap<char, ()>, HashMapStorage<char, ()>> =
            Network::new();

        net.add_nodes(('A'..='Z').map(|c| (c, ())));

        net.add_arcs(('A'..'C').zip('B'..='C').map(|(a, b)| (a, b, ())))
            .unwrap();
        net.add_arcs(('X'..'Z').zip('Y'..='Z').map(|(a, b)| (a, b, ())))
            .unwrap();

        net.add_arcs(('C'..'M').zip('D'..='M').map(|(a, b)| (a, b, ())))
            .unwrap();

        net.add_arcs(('N'..'X').zip('O'..='X').map(|(a, b)| (a, b, ())))
            .unwrap();

        net.add_arcs([('C', 'N', ()), ('M', 'X', ()), ('W', 'X', ())])
            .unwrap();

        let in_nodes = intermediate_nodes(&net, [&'C', &'W'], Direction::Undirected);

        let expected = Vec::from([
            &'C', &'D', &'E', &'F', &'G', &'H', &'I', &'J', &'K', &'L', &'M', &'N', &'O', &'P',
            &'Q', &'R', &'S', &'T', &'U', &'V', &'W', &'X',
        ]);

        let mut in_nodes: Vec<&char> = in_nodes.into_iter().collect();
        in_nodes.sort();

        assert_eq!(in_nodes, expected);
    }

    #[test]
    fn loop_intermediate_nodes() {
        let mut net: Network<char, (), (), HashMap<char, ()>, HashMapStorage<char, ()>> =
            Network::new();

        net.add_nodes(('A'..='D').map(|c| (c, ())));

        net.add_arcs([
            ('A', 'B', ()),
            ('A', 'C', ()),
            ('C', 'D', ()),
            ('B', 'D', ()),
        ])
        .unwrap();

        let in_nodes = intermediate_nodes(&net, [&'A', &'D'], Direction::Undirected);

        let expected = Vec::from([&'A', &'B', &'C', &'D']);

        let mut in_nodes: Vec<&char> = in_nodes.into_iter().collect();
        in_nodes.sort();

        assert_eq!(in_nodes, expected);
    }

    #[test]
    fn intermediate_nodes_misc() {
        let mut x_shaped: Network<char, (), (), HashMap<_, _>, HashMapStorage<_, _>> =
            Network::new();

        x_shaped.add_nodes(('A'..='E').map(|c| (c, ())));

        x_shaped
            .add_arcs([
                ('A', 'C', ()),
                ('B', 'C', ()),
                ('C', 'D', ()),
                ('C', 'E', ()),
            ])
            .unwrap();

        let in_nodes = intermediate_nodes(&x_shaped, [&'D', &'E'], Direction::Undirected);

        let expected = Vec::from([&'C', &'D', &'E']);

        let mut in_nodes: Vec<&char> = in_nodes.into_iter().collect();
        in_nodes.sort();

        assert_eq!(in_nodes, expected);
    }
}