neo-decompiler 0.10.1

use std::collections::BTreeMap;

use crate::decompiler::cfg::{BlockId, Cfg};

use super::traversal::{idom_parent, reverse_post_order};

/// Compute immediate dominators using the Cooper-Harvey-Kennedy algorithm.
///
/// For each block n, IDOM(n) is the unique block that:
/// - Strictly dominates n
/// - Does not strictly dominate any other block that dominates n
pub(super) fn compute_immediate_dominators(cfg: &Cfg) -> BTreeMap<BlockId, Option<BlockId>> {
    let mut idom: BTreeMap<BlockId, Option<BlockId>> = BTreeMap::new();
    let entry_id = cfg.entry_block().map(|b| b.id);

    for block in cfg.blocks() {
        let block_id = block.id;
        idom.insert(
            block_id,
            if Some(block_id) == entry_id {
                Some(block_id)
            } else {
                None
            },
        );
    }

    let rpo = reverse_post_order(cfg);
    let mut changed = true;
    let mut iteration_count = 0u32;
    while changed {
        iteration_count += 1;
        if iteration_count > 1000 {
            break;
        }
        changed = false;

        for &block_id in &rpo {
            if Some(block_id) == entry_id {
                continue;
            }

            let new_idom = intersect_dominators(cfg, block_id, &idom);
            let current_value = idom.get(&block_id).and_then(|o| *o);
            if current_value != new_idom {
                idom.insert(block_id, new_idom);
                changed = true;
            }
        }
    }

    idom
}

/// Build the dominator tree from immediate dominator relationships.
///
/// The dominator tree has edges from each block to its immediate dominator.
/// This creates a tree rooted at the entry block.
pub(super) fn build_dominator_tree(
    idom: &BTreeMap<BlockId, Option<BlockId>>,
) -> BTreeMap<BlockId, Vec<BlockId>> {
    let mut tree: BTreeMap<BlockId, Vec<BlockId>> = BTreeMap::new();

    for &block in idom.keys() {
        tree.entry(block).or_default();
    }

    for (&block, &opt_idom) in idom {
        if let Some(idom_block) = opt_idom {
            if idom_block != block {
                tree.entry(idom_block).or_default().push(block);
            }
        }
    }

    tree
}

/// Find the intersection of dominators for all predecessors of a block.
fn intersect_dominators(
    cfg: &Cfg,
    block: BlockId,
    idom: &BTreeMap<BlockId, Option<BlockId>>,
) -> Option<BlockId> {
    let predecessors = cfg.predecessors(block);

    if predecessors.is_empty() {
        return None;
    }

    let mut result = None;

    for pred in predecessors.iter() {
        let pred_idom = idom.get(pred).copied().flatten();

        result = match result {
            None => {
                if pred_idom.is_some() {
                    Some(*pred)
                } else {
                    None
                }
            }
            Some(current) => match pred_idom {
                None => Some(current),
                Some(_) => Some(find_common_dominator(current, *pred, idom)),
            },
        };
    }

    result
}

/// Find the least common ancestor (dominator) of two blocks.
fn find_common_dominator(
    mut finger1: BlockId,
    mut finger2: BlockId,
    idom: &BTreeMap<BlockId, Option<BlockId>>,
) -> BlockId {
    let mut depth1 = depth_in_dominator_tree(finger1, idom);
    let mut depth2 = depth_in_dominator_tree(finger2, idom);

    let mut iterations = 0;
    const MAX_ITERATIONS: usize = 1000;

    while depth1 > depth2 {
        let Some(next) = idom_parent(idom, finger1) else {
            return finger1;
        };
        finger1 = next;
        depth1 -= 1;
        iterations += 1;
        if iterations > MAX_ITERATIONS {
            return finger1;
        }
    }
    while depth2 > depth1 {
        let Some(next) = idom_parent(idom, finger2) else {
            return finger1;
        };
        finger2 = next;
        depth2 -= 1;
        iterations += 1;
        if iterations > MAX_ITERATIONS {
            return finger1;
        }
    }

    while finger1 != finger2 {
        let (Some(next1), Some(next2)) = (idom_parent(idom, finger1), idom_parent(idom, finger2))
        else {
            return finger1;
        };
        finger1 = next1;
        finger2 = next2;
        iterations += 1;
        if iterations > MAX_ITERATIONS {
            return finger1;
        }
    }

    finger1
}

fn depth_in_dominator_tree(block: BlockId, idom: &BTreeMap<BlockId, Option<BlockId>>) -> usize {
    let max_depth = idom.len();
    let mut depth = 1;
    let mut current = idom_parent(idom, block);

    while let Some(idom_block) = current {
        if depth >= max_depth {
            break;
        }
        depth += 1;
        current = idom_parent(idom, idom_block);
    }

    depth
}