libpitust 0.1.2

//! Module `scc_tree` constructs a tree from a graph such that each node
//! represents a subgraph of its parent. The root of this graph represents the
//! entire graph. Each non-root node is a strongly connected component within
//! the parent.

use std::collections::HashSet;

use crate::{
    graph::{GraphBase, GraphSucc, node_map::NodeMap},
    index::vec::IndexVec,
    index_type,
};

index_type!(pub TreeId);
index_type!(NodeId);

pub struct SCCTree<G: GraphBase> {
    node_tree: NodeMap<G, TreeId>,
    tree_parent: IndexVec<TreeId, Option<TreeId>>,
    tree_roots: IndexVec<TreeId, Vec<G::Node>>,
}
impl<G: GraphBase> SCCTree<G> {
    pub fn parent(&self, tree: TreeId) -> Option<TreeId> {
        self.tree_parent[tree]
    }
    pub fn tree_id(&self, node: G::Node) -> TreeId {
        self.node_tree[node]
    }
    pub fn roots(&self, id: TreeId) -> impl DoubleEndedIterator<Item = G::Node> {
        self.tree_roots[id].iter().copied()
    }
}

pub fn scc_tree<G: GraphBase + GraphSucc>(graph: G, roots: &[G::Node]) -> SCCTree<G> {
    // construct an interned version of the graph for the core algo
    let mut tree_map = IndexVec::<NodeId, G::Node>::new();
    let node_intern = NodeMap::new_with(graph, |node| tree_map.push(node));
    let mut tree_sp = tree_map.with_map(|_, _| (vec![], vec![]));
    let mut tree_once = tree_map.with_map(|_, _| false);
    let mut tree_assign = tree_map.with_map(|_, _| TreeId::ZERO);

    for node in graph.nodes() {
        let pred = node_intern[node];
        for succ in graph.succ_of(node) {
            let succ = node_intern[succ];
            tree_sp[succ].0.push(pred);
            tree_sp[pred].1.push(succ);
        }
    }

    let mut parent = IndexVec::new();
    let mut tree_roots = IndexVec::new();
    parent.push(TreeId::ZERO);
    let mut roots: Vec<NodeId> = roots.iter().map(|&node| node_intern[node]).collect();
    roots.sort_unstable();
    roots.dedup();
    tree_roots.push(roots);

    scc_tree_impl(
        &tree_sp,
        &mut parent,
        &mut tree_assign,
        &mut tree_once,
        &mut tree_roots,
        TreeId::ZERO,
    );

    let node_tree = NodeMap::new_with(graph, |node| tree_assign[node_intern[node]]);

    SCCTree {
        node_tree,
        tree_roots: tree_roots
            .with_map(|_, nodes| nodes.iter().copied().map(|n| tree_map[n]).collect()),
        tree_parent: parent.with_map(|this, &par| {
            if this == TreeId::ZERO {
                None
            } else {
                Some(par)
            }
        }),
    }
}

fn scc_tree_impl(
    graph: &IndexVec<NodeId, (Vec<NodeId>, Vec<NodeId>)>,
    parent: &mut IndexVec<TreeId, TreeId>,
    assign: &mut IndexVec<NodeId, TreeId>,
    once: &mut IndexVec<NodeId, bool>,
    tree_roots: &mut IndexVec<TreeId, Vec<NodeId>>,
    // roots: &mut IndexVec<NodeId, bool>,
    elem: TreeId,
) {
    // find roots of the SCC
    let mut roots = core::mem::take(&mut tree_roots[elem]);
    for (node, &id) in assign.iter() {
        if id != elem {
            continue;
        }
        for &pred in &graph[node].0 {
            if assign[pred] != elem {
                roots.push(node);
                break;
            }
        }
    }
    roots.sort_unstable();
    roots.dedup();
    assert!(roots.len() != 0, "visit elem {elem:?}");

    // forwards pass
    let mut order = vec![];
    for &root in &roots {
        once[root] = true;
    }
    for &root in &roots {
        scan_fwd(once, &mut order, assign, graph, root, elem);
    }
    // reset the onceset
    for (_, o) in &mut *once {
        *o = false;
    }
    for &root in &roots {
        once[root] = true;
    }

    // backwards pass
    let mut subwalk = vec![];
    let mut once2 = HashSet::new();
    for &node in order.iter().rev() {
        let mut elems = vec![];

        if once[node] {
            if roots.binary_search(&node).is_err() {
                continue;
            } else {
                elems.push(node);
            }
        } else {
            once[node] = true;
            scan_rev(once, &mut elems, assign, graph, node, elem);
        }

        let id = parent.push(elem);
        for &node in &elems {
            if !once2.insert(node) {
                panic!("bug: double-assigned {node:?}");
            }
            assert!(assign[node] == elem);
            assign[node] = id;
        }
        if elems.len() > 1 {
            tree_roots.push(vec![]);
            subwalk.push(id);
        } else {
            tree_roots.push(vec![node]);
        }
    }
    for (_, o) in &mut *once {
        *o = false;
    }

    tree_roots[elem] = roots;

    // compute SCC subtree
    for id in subwalk {
        scc_tree_impl(graph, parent, assign, once, tree_roots, id);
    }
}
fn scan_fwd(
    once: &mut IndexVec<NodeId, bool>,
    order: &mut Vec<NodeId>,
    assign: &IndexVec<NodeId, TreeId>,
    graph: &IndexVec<NodeId, (Vec<NodeId>, Vec<NodeId>)>,
    node: NodeId,
    tree: TreeId,
) {
    if assign[node] != tree {
        return;
    }
    once[node] = true;
    for &s in &graph[node].1 {
        if !once[s] {
            scan_fwd(once, order, assign, graph, s, tree);
        }
    }
    order.push(node);
}
fn scan_rev(
    once: &mut IndexVec<NodeId, bool>,
    order: &mut Vec<NodeId>,
    assign: &IndexVec<NodeId, TreeId>,
    graph: &IndexVec<NodeId, (Vec<NodeId>, Vec<NodeId>)>,
    node: NodeId,
    tree: TreeId,
) {
    if assign[node] != tree {
        return;
    }

    once[node] = true;
    for &s in &graph[node].0 {
        if !once[s] {
            scan_rev(once, order, assign, graph, s, tree);
        }
    }
    order.push(node);
}