Calculate domainance tree as fixed point

dom_tree.rs

pub trait AssocSet<X, Y: Clone + Default> {
    fn get(&self, target: X) -> Y;
    fn set(&mut self, key: X, val: Y);
}

pub trait DFSGraph {
    type Identifier: Clone + PartialEq + Eq + Copy;
    type Set<Y>: AssocSet<Self::Identifier, Y>
    where
        Y: Clone + Default;
    type SuccessorIter<'a>: Iterator<Item = Self::Identifier>
    where
        Self: 'a;

    fn create_set<Y>(&self) -> Self::Set<Y>
    where
        Y: Clone + Default;
    fn outgoing_edges<'a>(&'a self, id: Self::Identifier) -> Self::SuccessorIter<'a>;
    fn post_order_sequence(&self, root: Self::Identifier) -> Vec<Self::Identifier> {
        let mut stack = Vec::new();
        let mut visited = self.create_set();
        let mut visitor = |i| stack.push(i);
        self.post_order_visit(&mut visited, root, &mut visitor);
        return stack;
    }
    fn post_order_visit<F>(&self, set: &mut Self::Set<bool>, current: Self::Identifier, f: &mut F)
    where
        F: FnMut(Self::Identifier),
    {
        if set.get(current) {
            return;
        }

        set.set(current, true);

        for i in self.outgoing_edges(current) {
            self.post_order_visit(set, i, f);
        }

        f(current);
    }
}

pub trait DomTree: DFSGraph + Sized {
    type MutDomIter<'a>: Iterator<Item = &'a mut Option<Self::Identifier>>
    where
        Self: 'a;
    type PredecessorIter<'a>: Iterator<Item = Self::Identifier>
    where
        Self: 'a;

    fn dom(&self, id: Self::Identifier) -> Option<Self::Identifier>;
    fn set_dom(&mut self, id: Self::Identifier, target: Option<Self::Identifier>);
    fn predecessor_iter<'a>(&'a self, id: Self::Identifier) -> Self::PredecessorIter<'a>;
    fn doms_mut<'a>(&'a mut self) -> Self::MutDomIter<'a>;
    fn populate_dom(&mut self, root: Self::Identifier) {
        // two-pointer algorithm to trace LCA.
        fn intersect<D: DomTree>(
            tree: &D,
            order_map: &D::Set<usize>,
            mut x: D::Identifier,
            mut y: D::Identifier,
        ) -> D::Identifier {
            unsafe {
                while x != y {
                    while order_map.get(x) < order_map.get(y) {
                        // safe because we are processing in reverse post order.
                        x = tree.dom(x).unwrap_unchecked();
                    }
                    while order_map.get(y) < order_map.get(x) {
                        // safe because we are processing in reverse post order.
                        y = tree.dom(y).unwrap_unchecked();
                    }
                }
            }
            return x;
        }
        // initialize all immediate dominators to None.
        self.doms_mut().for_each(|x| *x = None);
        // set root to be its own immediate dominator.
        self.set_dom(root, Some(root));

        // establish post oder using DFS.
        let post_order = self.post_order_sequence(root);
        let mut post_order_map = self.create_set();
        for (i, k) in post_order.iter().cloned().enumerate() {
            post_order_map.set(k, i);
        }

        // Iterate until the fixed point is reached.
        let mut changed = true;
        while changed {
            changed = false;
            for (order, i) in post_order.iter().cloned().enumerate().rev() {
                if i == root {
                    continue;
                }
                let dom = unsafe {
                    self.predecessor_iter(i)
                        .filter(|x| post_order_map.get(*x) > order)
                        .next()
                        // safe because we are processing in reverse post order.
                        .unwrap_unchecked()
                };
                let dom = self
                    .predecessor_iter(i)
                    .filter(|x| *x != dom && self.dom(*x).is_some())
                    .fold(dom, |dom, x| intersect(self, &post_order_map, x, dom));

                if self.dom(i).map(|x| x != dom).unwrap_or(true) {
                    self.set_dom(i, Some(dom));
                    changed = true;
                }
            }
        }
    }
}