misc_iterators 0.3.3

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

/// Iterator that performs a classical DFS on a graph.
pub struct BFSIterator<V, E, I>
where
    V: Clone + Hash + Eq,
    E: FnMut(&V, usize, &V) -> I,
    I: Iterator<Item = V>,
{
    queue: VecDeque<(V, usize, V)>,
    seen_nodes: HashSet<V>,
    edges: E,
}

impl<'a, V, E, I> BFSIterator<V, E, I>
where
    V: Clone + Hash + Eq,
    E: FnMut(&V, usize, &V) -> I,
    I: Iterator<Item = V>,
{
    /// Constructs a DFS iterator that starts spanning the graph at `start`.
    /// `edges` is a function that returns an iterator over the neighbours
    /// of the specified node.
    pub fn new(start: V, edges: E) -> BFSIterator<V, E, I>
    {
        Self {
            queue: VecDeque::from(vec![(start.clone(), 0, start.clone())]),
            seen_nodes: HashSet::from_iter(vec![start]),
            edges,
        }
    }
}

impl<'a, V, E, I> Iterator for BFSIterator<V, E, I>
where
    V: Clone + Hash + Eq + Debug,
    E: FnMut(&V, usize, &V) -> I,
    I: Iterator<Item = V>,
{
    type Item = (V, usize, V);

    fn next(&mut self) -> Option<Self::Item>
    {
        if self.queue.is_empty()
        {
            return None;
        }

        let (vert, depth, ancestor) = self.queue.pop_front().unwrap();

        for n in (self.edges)(&vert, depth, &ancestor)
        {
            if !self.seen_nodes.contains(&n)
            {
                self.seen_nodes.insert(n.clone());
                self.queue.push_back((n, depth + 1, vert.clone()))
            }
        }

        Some((vert, depth, ancestor))
    }
}

// +| Tests |+ =======================================================
#[cfg(test)]
mod tests
{
    use std::collections::{HashMap, HashSet};

    use crate::breadth_first_search::BFSIterator;

    #[test]
    fn test_breadth_first_search()
    {
        let mut graph = HashMap::new();

        graph.insert(1, vec![2, 3]);
        graph.insert(2, vec![1, 4]);
        graph.insert(3, vec![1, 5]);
        graph.insert(4, vec![2]);
        graph.insert(5, vec![3]);

        let start = 1;
        let bfs =
            BFSIterator::new(start, |i, _, _| graph.get(i).unwrap().iter().map(|v| *v));

        let result: HashSet<_> = bfs.map(|(i, _, _)| i).collect();

        let expected_result: HashSet<usize> = vec![1, 2, 3, 4, 5].into_iter().collect();

        assert_eq!(result, expected_result);
    }
}