mapgraph/algo/

dom.rs

//! Contains an implementation of an algorithm to build a map of immediate dominators in a
//! [`Graph`](crate::Graph).

use crate::{
    graph::{iter::EdgesWalker, Edge, Edges, Node, Nodes},
    map::{IntKeyMap, Map},
    FrozenGraph,
};
use alloc::vec;
use core::fmt::Debug;

/// A trait for graph weights that store indices of their immediate dominators.
pub trait WithImmDom {
    /// The type for the immediate dominator index.
    type Ty: Copy + Eq + Debug + 'static;

    /// Returns the immediate dominator of a weight.
    fn imm_dom(&self) -> Option<Self::Ty>;

    /// Sets the immediate dominator of a weight.
    fn set_imm_dom(&mut self, dom: Self::Ty);
}

fn calc_imm_dominators_impl<N, E, NI, EI, NM, EM>(
    nodes: &mut Nodes<N, EI, NM>,
    edges: &Edges<E, NI, EI, EM>,
    node_index: NI,
    entry_idx: NI,
) where
    N: WithImmDom<Ty = NI>,
    NI: Copy + Eq + Debug + 'static,
    EI: Copy + Eq + Debug + 'static,
    NM: Map<Node<N, EI>, Key = NI>,
    EM: IntKeyMap<Edge<E, NI, EI>, Key = EI>,
{
    let mut successors = nodes.walk_successors(node_index);

    while let Some(successor_idx) = successors.walk_edges_next(edges) {
        // In case we reach the entry block this is definitely a back edge (since there is no
        // other way to reach the entry block twice).
        if successor_idx == entry_idx {
            //back_edges.insert(edge_idx);
            continue;
        }

        match nodes[successor_idx].weight().imm_dom() {
            None => {
                // In case a successor has no immediate dominator it will be us (until some
                // other predecessor of the block finds out).
                nodes[successor_idx].weight_mut().set_imm_dom(node_index);
                calc_imm_dominators_impl(nodes, edges, successor_idx, entry_idx);
            }
            Some(dominator) => {
                // We should never visit the same path.
                assert_ne!(dominator, node_index);

                // TODO: Maybe it's a good idea to cache current dominance path for each block?
                // Caching seems nice at first glance since it allows to avoid rebuilding the
                // dominance path, however it introduces new costs:
                // 1. The consumed memory.
                // 2. Each time a block is found that has wrong immediate dominator all blocks
                //    descending from it in the dominator tree must update their caches.
                let mut left_idx = dominator;
                let mut left_path = vec![dominator];
                let mut right_idx = node_index;
                let mut right_path = vec![node_index];
                loop {
                    // Left path is the path from the successor to the entry block. In case we
                    // find an intersection with the right path we set it as the new immediate
                    // dominator for the successor and exit the loop.
                    if let Some(left_dominator) = nodes[left_idx].weight_mut().imm_dom() {
                        if right_path.contains(&left_dominator) {
                            nodes[successor_idx]
                                .weight_mut()
                                .set_imm_dom(left_dominator);
                            break;
                        }
                        left_idx = left_dominator;
                        left_path.push(left_dominator);
                    }

                    // Right path is the path from the current block to the entry block. In case
                    // we find an intersection with the left path we set it as the new immediate
                    // dominator for the successor and exit the loop. In case we find the
                    // successor in the right path we have found a loop. Remember that info and
                    // exit the loop without modifying anything.
                    if let Some(right_dominator) = nodes[right_idx].weight_mut().imm_dom() {
                        if right_dominator == successor_idx {
                            //back_edges.insert(edge_idx);
                            break;
                        } else if left_path.contains(&right_dominator) {
                            nodes[successor_idx]
                                .weight_mut()
                                .set_imm_dom(right_dominator);
                            break;
                        }
                        right_idx = right_dominator;
                        right_path.push(right_dominator);
                    }
                }
            }
        }
    }
}

/// Calculates immediate dominators map for a graph.
///
/// The dominator indices are stored into node weights using the [`WithImmDom`] trait.
pub fn calc_imm_dominators<N, E, NI, EI, NM, EM>(
    graph: &mut FrozenGraph<N, E, NI, EI, NM, EM>,
    entry_index: NI,
) where
    N: WithImmDom<Ty = NI>,
    NI: Copy + Eq + Debug + 'static,
    EI: Copy + Eq + Debug + 'static,
    NM: Map<Node<N, EI>, Key = NI>,
    EM: IntKeyMap<Edge<E, NI, EI>, Key = EI>,
{
    let (nodes, edges) = graph.as_nodes_mut_and_edges();
    calc_imm_dominators_impl(nodes, edges, entry_index, entry_index)
}

#[cfg(all(test, feature = "slotmap"))]
mod tests {
    use super::*;
    use crate::{aliases::SlotMapGraph, map::slotmap::NodeIndex};

    struct TestNode(Option<NodeIndex>);

    impl Default for TestNode {
        fn default() -> Self {
            TestNode(None)
        }
    }

    impl WithImmDom for TestNode {
        type Ty = NodeIndex;

        fn imm_dom(&self) -> Option<Self::Ty> {
            self.0
        }

        fn set_imm_dom(&mut self, dom: Self::Ty) {
            self.0 = Some(dom);
        }
    }

    #[test]
    fn test_imm_dom_simple_diamond() {
        let mut graph = SlotMapGraph::<TestNode, ()>::with_capacities(4, 4);
        let entry_idx = graph.add_default_node();
        let left_idx = graph.add_default_node();
        let right_idx = graph.add_default_node();
        let end_idx = graph.add_default_node();

        graph.add_edge((), entry_idx, left_idx).unwrap();
        graph.add_edge((), entry_idx, right_idx).unwrap();
        graph.add_edge((), left_idx, end_idx).unwrap();
        graph.add_edge((), right_idx, end_idx).unwrap();

        calc_imm_dominators(&mut graph, entry_idx);

        assert_eq!(graph.node_weight(entry_idx).unwrap().imm_dom(), None);
        assert_eq!(
            graph.node_weight(left_idx).unwrap().imm_dom(),
            Some(entry_idx)
        );
        assert_eq!(
            graph.node_weight(right_idx).unwrap().imm_dom(),
            Some(entry_idx)
        );
        assert_eq!(
            graph.node_weight(end_idx).unwrap().imm_dom(),
            Some(entry_idx)
        );
    }

    #[test]
    fn test_imm_dom_simple_loop() {
        let mut graph = SlotMapGraph::<TestNode, ()>::with_capacities(5, 5);
        let entry_idx = graph.add_default_node();
        let loop_entry_idx = graph.add_default_node();
        let loop_body_idx = graph.add_default_node();
        let loop_exit_idx = graph.add_default_node();
        let end_idx = graph.add_default_node();

        graph.add_edge((), entry_idx, loop_entry_idx).unwrap();
        graph.add_edge((), loop_entry_idx, loop_body_idx).unwrap();
        graph.add_edge((), loop_body_idx, loop_exit_idx).unwrap();
        graph.add_edge((), loop_exit_idx, loop_entry_idx).unwrap();
        graph.add_edge((), loop_entry_idx, end_idx).unwrap();

        calc_imm_dominators(&mut graph, entry_idx);

        assert_eq!(graph.node_weight(entry_idx).unwrap().imm_dom(), None);
        assert_eq!(
            graph.node_weight(loop_entry_idx).unwrap().imm_dom(),
            Some(entry_idx)
        );
        assert_eq!(
            graph.node_weight(loop_body_idx).unwrap().imm_dom(),
            Some(loop_entry_idx)
        );
        assert_eq!(
            graph.node_weight(loop_exit_idx).unwrap().imm_dom(),
            Some(loop_body_idx)
        );
        assert_eq!(
            graph.node_weight(end_idx).unwrap().imm_dom(),
            Some(loop_entry_idx)
        );
    }

    #[test]
    fn test_imm_dom_nested_loop() {
        let mut graph = SlotMapGraph::<TestNode, ()>::with_capacities(9, 10);
        let entry_idx = graph.add_default_node();
        let outer_loop_entry_idx = graph.add_default_node();
        let outer_loop_start_idx = graph.add_default_node();
        let inner_loop_entry_idx = graph.add_default_node();
        let inner_loop_body_idx = graph.add_default_node();
        let inner_loop_exit_idx = graph.add_default_node();
        let outer_loop_end_idx = graph.add_default_node();
        let outer_loop_exit_idx = graph.add_default_node();
        let end_idx = graph.add_default_node();

        graph.add_edge((), entry_idx, outer_loop_entry_idx).unwrap();
        graph
            .add_edge((), outer_loop_entry_idx, outer_loop_start_idx)
            .unwrap();
        graph.add_edge((), outer_loop_entry_idx, end_idx).unwrap();
        graph
            .add_edge((), outer_loop_start_idx, inner_loop_entry_idx)
            .unwrap();
        graph
            .add_edge((), inner_loop_entry_idx, inner_loop_body_idx)
            .unwrap();
        graph
            .add_edge((), inner_loop_entry_idx, outer_loop_end_idx)
            .unwrap();
        graph
            .add_edge((), inner_loop_body_idx, inner_loop_exit_idx)
            .unwrap();
        graph
            .add_edge((), inner_loop_exit_idx, inner_loop_entry_idx)
            .unwrap();
        graph
            .add_edge((), outer_loop_end_idx, outer_loop_exit_idx)
            .unwrap();
        graph
            .add_edge((), outer_loop_exit_idx, outer_loop_entry_idx)
            .unwrap();

        calc_imm_dominators(&mut graph, entry_idx);

        assert_eq!(graph.node_weight(entry_idx).unwrap().imm_dom(), None);
        assert_eq!(
            graph.node_weight(outer_loop_entry_idx).unwrap().imm_dom(),
            Some(entry_idx)
        );
        assert_eq!(
            graph.node_weight(outer_loop_start_idx).unwrap().imm_dom(),
            Some(outer_loop_entry_idx)
        );
        assert_eq!(
            graph.node_weight(inner_loop_entry_idx).unwrap().imm_dom(),
            Some(outer_loop_start_idx)
        );
        assert_eq!(
            graph.node_weight(inner_loop_body_idx).unwrap().imm_dom(),
            Some(inner_loop_entry_idx)
        );
        assert_eq!(
            graph.node_weight(inner_loop_exit_idx).unwrap().imm_dom(),
            Some(inner_loop_body_idx)
        );
        assert_eq!(
            graph.node_weight(outer_loop_end_idx).unwrap().imm_dom(),
            Some(inner_loop_entry_idx)
        );
        assert_eq!(
            graph.node_weight(outer_loop_exit_idx).unwrap().imm_dom(),
            Some(outer_loop_end_idx)
        );
        assert_eq!(
            graph.node_weight(end_idx).unwrap().imm_dom(),
            Some(outer_loop_entry_idx)
        );
    }
}