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

use crate::{
    graph::{Edge, Node},
    map::{ExtKeyMap, IntKeyMap, Map},
    FrozenGraph,
};
use alloc::vec;
use core::fmt::Debug;

fn calc_imm_dominators_impl<N, E, NI, EI, NM, EM, SM>(
    graph: &FrozenGraph<N, E, NI, EI, NM, EM>,
    node_index: NI,
    entry_idx: NI,
    imm_dominators: &mut SM,
) where
    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>,
    SM: ExtKeyMap<NI, Key = NI>,
{
    let successors = graph.successors(node_index);

    for (successor_idx, _) in successors {
        // 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 imm_dominators.get_mut(successor_idx) {
            None => {
                // In case a successor has no immediate dominator it will be us (until some
                // other predecessor of the block finds out).
                imm_dominators.replace(successor_idx, node_index);
                calc_imm_dominators_impl(graph, successor_idx, entry_idx, imm_dominators);
            }
            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];
                let dominator_ptr = dominator as *mut NI;
                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) = imm_dominators.get(left_idx) {
                        if right_path.contains(&left_dominator) {
                            unsafe { *dominator_ptr = 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) = imm_dominators.get(right_idx) {
                        if right_dominator == successor_idx {
                            //back_edges.insert(edge_idx);
                            break;
                        } else if left_path.contains(&right_dominator) {
                            unsafe { *dominator_ptr = right_dominator };
                            break;
                        }
                        right_idx = right_dominator;
                        right_path.push(right_dominator);
                    }
                }
            }
        }
    }
}

/// Calculates immediate dominators map for a graph.
pub fn calc_imm_dominators<N, E, NI, EI, NM, EM, SM>(
    graph: &FrozenGraph<N, E, NI, EI, NM, EM>,
    entry_index: NI,
    imm_dominators: &mut SM,
) where
    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>,
    SM: ExtKeyMap<NI, Key = NI>,
{
    calc_imm_dominators_impl(graph, entry_index, entry_index, imm_dominators)
}

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

    #[test]
    fn test_imm_dom_simple_diamond() {
        let mut graph = SlotMapGraph::<(), ()>::with_capacities(4, 4);
        let entry_idx = graph.add_node(());
        let left_idx = graph.add_node(());
        let right_idx = graph.add_node(());
        let end_idx = graph.add_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();

        let mut imm_dom = SecondaryMap::with_capacity(4);
        calc_imm_dominators(&graph, entry_idx, &mut imm_dom);

        assert_eq!(imm_dom.get(entry_idx), None);
        assert_eq!(imm_dom.get(left_idx).copied(), Some(entry_idx));
        assert_eq!(imm_dom.get(right_idx).copied(), Some(entry_idx));
        assert_eq!(imm_dom.get(end_idx).copied(), Some(entry_idx));
    }

    #[test]
    fn test_imm_dom_simple_loop() {
        let mut graph = SlotMapGraph::<(), ()>::with_capacities(5, 5);
        let entry_idx = graph.add_node(());
        let loop_entry_idx = graph.add_node(());
        let loop_body_idx = graph.add_node(());
        let loop_exit_idx = graph.add_node(());
        let end_idx = graph.add_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();

        let mut imm_dom = SecondaryMap::with_capacity(4);
        calc_imm_dominators(&graph, entry_idx, &mut imm_dom);

        assert_eq!(imm_dom.get(entry_idx), None);
        assert_eq!(imm_dom.get(loop_entry_idx).copied(), Some(entry_idx));
        assert_eq!(imm_dom.get(loop_body_idx).copied(), Some(loop_entry_idx));
        assert_eq!(imm_dom.get(loop_exit_idx).copied(), Some(loop_body_idx));
        assert_eq!(imm_dom.get(end_idx).copied(), Some(loop_entry_idx));
    }

    #[test]
    fn test_imm_dom_nested_loop() {
        let mut graph = SlotMapGraph::<(), ()>::with_capacities(9, 10);
        let entry_idx = graph.add_node(());
        let outer_loop_entry_idx = graph.add_node(());
        let outer_loop_start_idx = graph.add_node(());
        let inner_loop_entry_idx = graph.add_node(());
        let inner_loop_body_idx = graph.add_node(());
        let inner_loop_exit_idx = graph.add_node(());
        let outer_loop_end_idx = graph.add_node(());
        let outer_loop_exit_idx = graph.add_node(());
        let end_idx = graph.add_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();

        let mut imm_dom = SecondaryMap::with_capacity(4);
        calc_imm_dominators(&graph, entry_idx, &mut imm_dom);

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