libpitust 0.1.0

//! Module `domtree` computes the dominator tree of a graph.

use crate::graph::{GraphBase, GraphPred, GraphRooted, GraphSucc, node_map::NodeMap};

struct BlockState<G: GraphBase> {
    /// Zero if the block is unreachable, and a nonzero number representing the
    /// index otherwise.
    postorder_idx: u32,

    /// The immediate dominator of the node
    immediate_dominator: Option<G::Node>,
}
impl<G: GraphBase> Default for BlockState<G> {
    fn default() -> Self {
        BlockState {
            postorder_idx: 0,
            immediate_dominator: None,
        }
    }
}

pub struct DominatorTree<G: GraphRooted + GraphPred + GraphSucc> {
    graph: G,
    block_state: NodeMap<G, BlockState<G>>,
    postorder: Vec<G::Node>,
    root_info: G::RootInfo,
}
impl<G: GraphRooted + GraphPred + GraphSucc> DominatorTree<G> {
    #[inline(always)]
    pub fn new(graph: G) -> DominatorTree<G> {
        let mut domtree = DominatorTree {
            graph,
            block_state: NodeMap::new(graph),
            postorder: Vec::new(),
            root_info: graph.root_node_info(),
        };
        domtree.compute_postorder();
        domtree.compute_dominator();

        debug_assert!(domtree.postorder.len() <= domtree.graph.nodes().count());
        domtree
    }

    #[inline(always)]
    pub fn is_reachable(&self, node: G::Node) -> bool {
        self.block_state[node].postorder_idx != 0
    }

    fn do_compute_postorder(&mut self, node: G::Node) {
        if self.block_state[node].postorder_idx != 0 {
            return;
        }

        self.block_state[node].postorder_idx = u32::MAX;
        for succ in self.graph.succ_of(node) {
            self.do_compute_postorder(succ);
        }
        self.postorder.push(node);
        self.block_state[node].postorder_idx = self.postorder.len() as u32;
    }
    fn compute_postorder(&mut self) {
        self.do_compute_postorder(self.graph.root_node());
    }

    pub fn common_dominator(&self, mut n1: G::Node, mut n2: G::Node) -> G::Node {
        // FIXME: this is an impressively bad implementation
        let mut walk1 = vec![n1];
        while let Some(prev) = self.block_state[n1].immediate_dominator {
            walk1.push(prev);
            n1 = prev;
        }

        if walk1.contains(&n2) {
            return n2;
        }
        while let Some(prev) = self.block_state[n2].immediate_dominator {
            if walk1.contains(&prev) {
                return prev;
            }
            n2 = prev;
        }

        unreachable!()
    }
    pub fn dominates(&self, node: G::Node, other: G::Node) -> bool {
        let mut other = other;
        loop {
            if other == node {
                return true;
            }
            if let Some(prev) = self.block_state[other].immediate_dominator {
                other = prev;
            } else {
                return false;
            }
        }
    }
    pub fn immediate_dominator(&self, node: G::Node) -> G::Node {
        if node == G::root_node_from_info(self.root_info) {
            return node;
        }

        self.block_state[node].immediate_dominator.unwrap()
    }
    pub fn iter_rpo(&self) -> impl DoubleEndedIterator<Item = G::Node> {
        self.postorder.iter().copied().rev()
    }

    fn compute_dominator(&mut self) {
        // FIXME: use a less shit dominator tree algo when perf becomes a problem

        // choose a random predecessor which "could" be our dominator
        for node in self.graph.nodes() {
            if node != G::root_node_from_info(self.root_info) && self.is_reachable(node) {
                let pred = self
                    .graph
                    .pred_of(node)
                    .find(|&n| {
                        self.block_state[n].postorder_idx > self.block_state[node].postorder_idx
                    })
                    .unwrap();
                self.block_state[node].immediate_dominator = Some(pred);
            }
        }

        loop {
            let mut changed = false;
            for node in self.graph.nodes() {
                if node == G::root_node_from_info(self.root_info) {
                    continue;
                }

                if !self.is_reachable(node) {
                    continue;
                }

                let mut idom = self.block_state[node].immediate_dominator.unwrap();
                let old_idom = idom;

                for pred in self.graph.pred_of(node) {
                    if !self.is_reachable(pred) {
                        continue;
                    }

                    idom = self.common_dominator(idom, pred);
                }
                if old_idom != idom {
                    changed = true;
                    self.block_state[node].immediate_dominator = Some(idom);
                }
            }

            if !changed {
                break;
            }
        }
    }
}