src/util/graph.rs

use std::borrow::Borrow;
use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::fmt;

pub struct Graph<N: Clone, E: Clone> {
    nodes: im_rc::OrdMap<N, im_rc::OrdMap<N, E>>,
}

impl<N: Eq + Ord + Clone, E: Default + Clone> Graph<N, E> {
    pub fn new() -> Graph<N, E> {
        Graph {
            nodes: im_rc::OrdMap::new(),
        }
    }

    pub fn add(&mut self, node: N) {
        self.nodes.entry(node).or_insert_with(im_rc::OrdMap::new);
    }

    pub fn link(&mut self, node: N, child: N) -> &mut E {
        self.nodes
            .entry(node)
            .or_insert_with(im_rc::OrdMap::new)
            .entry(child)
            .or_default()
    }

    /// Returns the graph obtained by reversing all edges.
    pub fn reversed(&self) -> Graph<N, E> {
        let mut ret = Graph::new();

        for n in self.iter() {
            ret.add(n.clone());
            for (m, e) in self.edges(n) {
                *ret.link(m.clone(), n.clone()) = e.clone();
            }
        }

        ret
    }

    pub fn contains<Q: ?Sized>(&self, k: &Q) -> bool
    where
        N: Borrow<Q>,
        Q: Ord + Eq,
    {
        self.nodes.contains_key(k)
    }

    pub fn edge(&self, from: &N, to: &N) -> Option<&E> {
        self.nodes.get(from)?.get(to)
    }

    pub fn edges(&self, from: &N) -> impl Iterator<Item = (&N, &E)> + use<'_, N, E> {
        self.nodes.get(from).into_iter().flat_map(|x| x.iter())
    }

    /// A topological sort of the `Graph`
    pub fn sort(&self) -> Vec<N> {
        let mut ret = Vec::new();
        let mut marks = BTreeSet::new();

        for node in self.nodes.keys() {
            self.sort_inner_visit(node, &mut ret, &mut marks);
        }

        ret
    }

    fn sort_inner_visit(&self, node: &N, dst: &mut Vec<N>, marks: &mut BTreeSet<N>) {
        if !marks.insert(node.clone()) {
            return;
        }

        for child in self.nodes[node].keys() {
            self.sort_inner_visit(child, dst, marks);
        }

        dst.push(node.clone());
    }

    pub fn iter(&self) -> impl Iterator<Item = &N> {
        self.nodes.keys()
    }

    pub fn len(&self) -> usize {
        self.nodes.len()
    }

    /// Checks if there is a path from `from` to `to`.
    pub fn is_path_from_to<'a>(&'a self, from: &'a N, to: &'a N) -> bool {
        let mut stack = vec![from];
        let mut seen = BTreeSet::new();
        seen.insert(from);
        while let Some(iter) = stack.pop().and_then(|p| self.nodes.get(p)) {
            for p in iter.keys() {
                if p == to {
                    return true;
                }
                if seen.insert(p) {
                    stack.push(p);
                }
            }
        }
        false
    }

    /// Resolves one of the paths from the given dependent package down to a leaf.
    ///
    /// The path return will be the shortest path, or more accurately one of the paths with the shortest length.
    ///
    /// Each element contains a node along with an edge except the first one.
    /// The representation would look like:
    ///
    /// (Node0,) -> (Node1, Edge01) -> (Node2, Edge12)...
    pub fn path_to_bottom<'a>(&'a self, pkg: &'a N) -> Vec<(&'a N, Option<&'a E>)> {
        self.path_to(pkg, |s, p| s.edges(p))
    }

    /// Resolves one of the paths from the given dependent package up to the root.
    ///
    /// The path return will be the shortest path, or more accurately one of the paths with the shortest length.
    ///
    /// Each element contains a node along with an edge except the first one.
    /// The representation would look like:
    ///
    /// (Node0,) -> (Node1, Edge01) -> (Node2, Edge12)...
    pub fn path_to_top<'a>(&'a self, pkg: &'a N) -> Vec<(&'a N, Option<&'a E>)> {
        self.path_to(pkg, |s, pk| {
            // Note that this implementation isn't the most robust per se, we'll
            // likely have to tweak this over time. For now though it works for what
            // it's used for!
            s.nodes
                .iter()
                .filter_map(|(p, adjacent)| adjacent.get(pk).map(|e| (p, e)))
        })
    }
}

impl<'s, N: Eq + Ord + Clone + 's, E: Default + Clone + 's> Graph<N, E> {
    fn path_to<'a, F, I>(&'s self, pkg: &'a N, fn_edge: F) -> Vec<(&'a N, Option<&'a E>)>
    where
        I: Iterator<Item = (&'a N, &'a E)>,
        F: Fn(&'s Self, &'a N) -> I,
        'a: 's,
    {
        let mut back_link = BTreeMap::new();
        let mut queue = VecDeque::from([pkg]);
        let mut last = pkg;

        while let Some(p) = queue.pop_front() {
            last = p;
            let mut out_edges = true;
            for (child, edge) in fn_edge(&self, p) {
                out_edges = false;
                back_link.entry(child).or_insert_with(|| {
                    queue.push_back(child);
                    (p, edge)
                });
            }
            if out_edges {
                break;
            }
        }

        let mut result = Vec::new();
        let mut next = last;
        while let Some((p, e)) = back_link.remove(&next) {
            result.push((next, Some(e)));
            next = p;
        }
        if result.iter().all(|(n, _)| n != &next) {
            result.push((next, None));
        }
        result.reverse();
        #[cfg(debug_assertions)]
        {
            for x in result.windows(2) {
                let [(n2, _), (n1, Some(e12))] = x else {
                    unreachable!()
                };
                assert!(std::ptr::eq(
                    self.edge(n1, n2).or(self.edge(n2, n1)).unwrap(),
                    *e12
                ));
            }
            let last = result.last().unwrap().0;
            let set: Vec<_> = result.iter().map(|(k, _)| k).collect();
            if !fn_edge(&self, last)
                .filter(|(e, _)| !set.contains(&e))
                .next()
                .is_none()
            {
                self.print_for_test();
                unreachable!("The last element in the path should not have outgoing edges");
            }
        }
        result
    }
}

#[test]
fn path_to_case() {
    let mut new = Graph::new();
    new.link(0, 3);
    new.link(1, 0);
    new.link(2, 0);
    new.link(2, 1);
    assert_eq!(
        new.path_to_bottom(&2),
        vec![(&2, None), (&0, Some(&())), (&3, Some(&()))]
    );
}

#[test]
fn path_to_self() {
    // Extracted from #12941
    let mut new: Graph<i32, ()> = Graph::new();
    new.link(0, 0);
    assert_eq!(new.path_to_bottom(&0), vec![(&0, Some(&()))]);
}

#[test]
fn reverse() {
    let mut new: Graph<i32, ()> = Graph::new();
    new.link(0, 1);
    new.link(0, 2);

    let mut expected: Graph<i32, ()> = Graph::new();
    expected.add(0);
    expected.link(1, 0);
    expected.link(2, 0);
    assert_eq!(new.reversed(), expected);
}

impl<N: Eq + Ord + Clone, E: Default + Clone> Default for Graph<N, E> {
    fn default() -> Graph<N, E> {
        Graph::new()
    }
}

impl<N: fmt::Display + Eq + Ord + Clone, E: Clone> fmt::Debug for Graph<N, E> {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(fmt, "Graph {{")?;

        for (n, e) in &self.nodes {
            writeln!(fmt, "  - {}", n)?;

            for n in e.keys() {
                writeln!(fmt, "    - {}", n)?;
            }
        }

        write!(fmt, "}}")?;

        Ok(())
    }
}

impl<N: Eq + Ord + Clone, E: Clone> Graph<N, E> {
    /// Prints the graph for constructing unit tests.
    ///
    /// For purposes of graph traversal algorithms the edge values do not matter,
    /// and the only value of the node we care about is the order it gets compared in.
    /// This constructs a graph with the same topology but with integer keys and unit edges.
    #[cfg(debug_assertions)]
    #[allow(clippy::print_stderr)]
    fn print_for_test(&self) {
        // Isolate and print a test case.
        let names = self
            .nodes
            .keys()
            .chain(self.nodes.values().flat_map(|vs| vs.keys()))
            .collect::<BTreeSet<_>>()
            .into_iter()
            .collect::<Vec<_>>();
        let mut new = Graph::new();
        for n1 in self.nodes.keys() {
            let name1 = names.binary_search(&n1).unwrap();
            new.add(name1);
            for n2 in self.nodes[n1].keys() {
                let name2 = names.binary_search(&n2).unwrap();
                *new.link(name1, name2) = ();
            }
        }
        eprintln!("Graph for tests = {new:#?}");
    }
}

impl<N: Eq + Ord + Clone, E: Eq + Clone> PartialEq for Graph<N, E> {
    fn eq(&self, other: &Graph<N, E>) -> bool {
        self.nodes.eq(&other.nodes)
    }
}
impl<N: Eq + Ord + Clone, E: Eq + Clone> Eq for Graph<N, E> {}

impl<N: Eq + Ord + Clone, E: Clone> Clone for Graph<N, E> {
    fn clone(&self) -> Graph<N, E> {
        Graph {
            nodes: self.nodes.clone(),
        }
    }
}