weirflow 0.1.0

use crate::ssa::{
    compute_dominance_frontiers, compute_dominators, place_phi_nodes, rename_variables,
    try_compute_dominators, try_compute_dominators_cpu, try_compute_dominators_detailed, Block,
    Cfg, DominatorPath, DOMINATOR_GPU_THRESHOLD,
};
use std::collections::HashMap;

#[test]
fn cytron_paper_fig_1() {
    let mut cfg = Cfg {
        entry: 1,
        blocks: HashMap::new(),
    };

    // Fig 1 CFG
    cfg.blocks.insert(
        1,
        Block {
            id: 1,
            preds: vec![],
            succs: vec![2],
            defs: [1].into(),
            uses: [].into(),
        },
    );
    cfg.blocks.insert(
        2,
        Block {
            id: 2,
            preds: vec![1, 3],
            succs: vec![3, 4],
            defs: [].into(),
            uses: [1].into(),
        },
    );
    cfg.blocks.insert(
        3,
        Block {
            id: 3,
            preds: vec![2],
            succs: vec![2],
            defs: [1].into(),
            uses: [].into(),
        },
    );
    cfg.blocks.insert(
        4,
        Block {
            id: 4,
            preds: vec![2],
            succs: vec![],
            defs: [].into(),
            uses: [1].into(),
        },
    );

    let doms = compute_dominators(&cfg).unwrap();
    let df = compute_dominance_frontiers(&cfg, &doms);
    let phi_nodes = place_phi_nodes(&cfg, &df);
    let _ssa = rename_variables(&cfg, &doms, &phi_nodes);

    assert!(phi_nodes.get(&2).unwrap().contains(&1));
}

#[test]
fn linear_chain() {
    let mut cfg = Cfg {
        entry: 1,
        blocks: HashMap::new(),
    };
    cfg.blocks.insert(
        1,
        Block {
            id: 1,
            preds: vec![],
            succs: vec![2],
            defs: [1].into(),
            uses: [].into(),
        },
    );
    cfg.blocks.insert(
        2,
        Block {
            id: 2,
            preds: vec![1],
            succs: vec![3],
            defs: [].into(),
            uses: [1].into(),
        },
    );
    cfg.blocks.insert(
        3,
        Block {
            id: 3,
            preds: vec![2],
            succs: vec![],
            defs: [1].into(),
            uses: [].into(),
        },
    );

    let doms = compute_dominators(&cfg).unwrap();
    let df = compute_dominance_frontiers(&cfg, &doms);
    let phi_nodes = place_phi_nodes(&cfg, &df);

    assert!(phi_nodes.is_empty() || phi_nodes.values().all(|v| v.is_empty()));
}

#[test]
fn diamond_with_phi() {
    let mut cfg = Cfg {
        entry: 1,
        blocks: HashMap::new(),
    };
    cfg.blocks.insert(
        1,
        Block {
            id: 1,
            preds: vec![],
            succs: vec![2, 3],
            defs: [].into(),
            uses: [].into(),
        },
    );
    cfg.blocks.insert(
        2,
        Block {
            id: 2,
            preds: vec![1],
            succs: vec![4],
            defs: [1].into(),
            uses: [].into(),
        },
    );
    cfg.blocks.insert(
        3,
        Block {
            id: 3,
            preds: vec![1],
            succs: vec![4],
            defs: [1].into(),
            uses: [].into(),
        },
    );
    cfg.blocks.insert(
        4,
        Block {
            id: 4,
            preds: vec![2, 3],
            succs: vec![],
            defs: [].into(),
            uses: [1].into(),
        },
    );

    let doms = compute_dominators(&cfg).unwrap();
    let df = compute_dominance_frontiers(&cfg, &doms);
    let phi_nodes = place_phi_nodes(&cfg, &df);

    assert!(phi_nodes.get(&4).unwrap().contains(&1));
}

#[test]
fn nested_loop_phi() {
    let mut cfg = Cfg {
        entry: 1,
        blocks: HashMap::new(),
    };
    cfg.blocks.insert(
        1,
        Block {
            id: 1,
            preds: vec![],
            succs: vec![2],
            defs: [1].into(),
            uses: [].into(),
        },
    );
    cfg.blocks.insert(
        2,
        Block {
            id: 2,
            preds: vec![1, 5],
            succs: vec![3, 6],
            defs: [].into(),
            uses: [].into(),
        },
    );
    cfg.blocks.insert(
        3,
        Block {
            id: 3,
            preds: vec![2, 4],
            succs: vec![4, 5],
            defs: [].into(),
            uses: [1].into(),
        },
    );
    cfg.blocks.insert(
        4,
        Block {
            id: 4,
            preds: vec![3],
            succs: vec![3],
            defs: [1].into(),
            uses: [].into(),
        },
    );
    cfg.blocks.insert(
        5,
        Block {
            id: 5,
            preds: vec![3],
            succs: vec![2],
            defs: [1].into(),
            uses: [].into(),
        },
    );
    cfg.blocks.insert(
        6,
        Block {
            id: 6,
            preds: vec![2],
            succs: vec![],
            defs: [].into(),
            uses: [1].into(),
        },
    );

    let doms = compute_dominators(&cfg).unwrap();
    let df = compute_dominance_frontiers(&cfg, &doms);
    let phi_nodes = place_phi_nodes(&cfg, &df);

    assert!(phi_nodes.get(&2).unwrap().contains(&1));
    assert!(phi_nodes.get(&3).unwrap().contains(&1));
}

#[test]
fn irreducible_loop() {
    let mut cfg = Cfg {
        entry: 1,
        blocks: HashMap::new(),
    };
    cfg.blocks.insert(
        1,
        Block {
            id: 1,
            preds: vec![],
            succs: vec![2, 3],
            defs: [1].into(),
            uses: [].into(),
        },
    );
    cfg.blocks.insert(
        2,
        Block {
            id: 2,
            preds: vec![1, 3],
            succs: vec![3],
            defs: [1].into(),
            uses: [].into(),
        },
    );
    cfg.blocks.insert(
        3,
        Block {
            id: 3,
            preds: vec![1, 2],
            succs: vec![2],
            defs: [].into(),
            uses: [1].into(),
        },
    );

    let doms = compute_dominators(&cfg).unwrap();
    let df = compute_dominance_frontiers(&cfg, &doms);
    let phi_nodes = place_phi_nodes(&cfg, &df);

    // An irreducible loop might still be processed by standard SSA, but usually requires node splitting or yields strange doms.
    // In our simplified CHK algorithm, it should complete. We just ensure it doesn't panic.
    assert!(phi_nodes.contains_key(&2) || phi_nodes.contains_key(&3) || true);
}

#[test]
fn try_compute_dominators_rejects_missing_entry() {
    let cfg = Cfg {
        entry: 99,
        blocks: HashMap::new(),
    };
    let err = crate::ssa::try_compute_dominators(&cfg).expect_err("missing entry");
    assert!(err.contains("entry block not found"), "{err}");
}

#[test]
fn try_compute_dominators_rejects_disconnected_entry() {
    let mut blocks = HashMap::new();
    blocks.insert(
        1,
        Block {
            id: 1,
            preds: vec![],
            succs: vec![],
            defs: [].into(),
            uses: [].into(),
        },
    );
    let cfg = Cfg { entry: 1, blocks };
    let doms = crate::ssa::try_compute_dominators(&cfg).unwrap();
    assert_eq!(doms.get(&1), Some(&1));
}

#[test]
fn try_compute_dominance_frontiers_empty_cfg() {
    let mut blocks = HashMap::new();
    blocks.insert(
        1,
        Block {
            id: 1,
            preds: vec![],
            succs: vec![],
            defs: [].into(),
            uses: [].into(),
        },
    );
    let cfg = Cfg { entry: 1, blocks };
    let doms = crate::ssa::try_compute_dominators(&cfg).unwrap();
    let df = crate::ssa::try_compute_dominance_frontiers(&cfg, &doms).unwrap();
    assert!(df.is_empty() || df.values().all(|v| v.is_empty()));
}

#[test]
fn try_place_phi_nodes_empty_df() {
    let mut blocks = HashMap::new();
    blocks.insert(
        1,
        Block {
            id: 1,
            preds: vec![],
            succs: vec![],
            defs: [].into(),
            uses: [].into(),
        },
    );
    let cfg = Cfg { entry: 1, blocks };
    let doms = crate::ssa::try_compute_dominators(&cfg).unwrap();
    let df = crate::ssa::try_compute_dominance_frontiers(&cfg, &doms).unwrap();
    let phi = crate::ssa::try_place_phi_nodes(&cfg, &df).unwrap();
    assert!(phi.is_empty());
}

#[test]
fn try_rename_variables_empty() {
    let mut blocks = HashMap::new();
    blocks.insert(
        1,
        Block {
            id: 1,
            preds: vec![],
            succs: vec![],
            defs: [].into(),
            uses: [].into(),
        },
    );
    let cfg = Cfg { entry: 1, blocks };
    let doms = crate::ssa::try_compute_dominators(&cfg).unwrap();
    let df = crate::ssa::try_compute_dominance_frontiers(&cfg, &doms).unwrap();
    let phi = crate::ssa::try_place_phi_nodes(&cfg, &df).unwrap();
    let ssa = crate::ssa::try_rename_variables(&cfg, &doms, &phi).unwrap();
    assert!(ssa.phi_nodes.is_empty());
    assert!(ssa.renamed_usages.is_empty());
}

#[cfg(test)]
mod dominator_threshold_tests {
    use super::{
        try_compute_dominators, try_compute_dominators_cpu, try_compute_dominators_detailed, Block,
        Cfg, DominatorPath, DOMINATOR_GPU_THRESHOLD,
    };
    use std::collections::{HashMap, HashSet};

    fn linear_cfg(n: u32) -> Cfg {
        let mut blocks = HashMap::with_capacity(n as usize);
        for i in 0..n {
            let preds = if i == 0 { vec![] } else { vec![i - 1] };
            let succs = if i + 1 < n { vec![i + 1] } else { vec![] };
            let mut defs = HashSet::new();
            defs.insert(i);
            blocks.insert(
                i,
                Block {
                    id: i,
                    preds,
                    succs,
                    defs,
                    uses: HashSet::new(),
                },
            );
        }
        Cfg { entry: 0, blocks }
    }

    #[test]
    fn threshold_exactly_1000_uses_cpu() {
        let cfg = linear_cfg(DOMINATOR_GPU_THRESHOLD);
        let (doms, path) = try_compute_dominators_detailed(&cfg).unwrap();
        assert_eq!(path, DominatorPath::Cpu);
        assert_eq!(doms.get(&0), Some(&0));
        assert_eq!(
            doms.get(&(DOMINATOR_GPU_THRESHOLD - 1)),
            Some(&(DOMINATOR_GPU_THRESHOLD - 2))
        );
    }

    #[test]
    fn threshold_1001_attempts_gpu() {
        let cfg = linear_cfg(DOMINATOR_GPU_THRESHOLD + 1);
        let (doms, path) = try_compute_dominators_detailed(&cfg).unwrap();
        assert!(matches!(
            path,
            DominatorPath::Gpu | DominatorPath::GpuFallback
        ));
        assert_eq!(doms.get(&0), Some(&0));
        assert_eq!(
            doms.get(&DOMINATOR_GPU_THRESHOLD),
            Some(&(DOMINATOR_GPU_THRESHOLD - 1))
        );
    }

    #[test]
    fn cpu_and_gpu_paths_agree_on_linear_chain() {
        let cfg = linear_cfg(DOMINATOR_GPU_THRESHOLD + 5);
        let cpu_doms = try_compute_dominators_cpu(&cfg).unwrap();
        let routed_doms = try_compute_dominators(&cfg).unwrap();
        assert_eq!(cpu_doms, routed_doms);
    }

    #[test]
    fn cpu_and_gpu_paths_agree_on_diamond() {
        let mut blocks = HashMap::new();
        blocks.insert(
            0,
            Block {
                id: 0,
                preds: vec![],
                succs: vec![1, 2],
                defs: HashSet::new(),
                uses: HashSet::new(),
            },
        );
        blocks.insert(
            1,
            Block {
                id: 1,
                preds: vec![0],
                succs: vec![3],
                defs: HashSet::new(),
                uses: HashSet::new(),
            },
        );
        blocks.insert(
            2,
            Block {
                id: 2,
                preds: vec![0],
                succs: vec![3],
                defs: HashSet::new(),
                uses: HashSet::new(),
            },
        );
        blocks.insert(
            3,
            Block {
                id: 3,
                preds: vec![1, 2],
                succs: vec![],
                defs: HashSet::new(),
                uses: HashSet::new(),
            },
        );
        let cfg = Cfg { entry: 0, blocks };
        let cpu_doms = try_compute_dominators_cpu(&cfg).unwrap();
        let routed_doms = try_compute_dominators(&cfg).unwrap();
        assert_eq!(cpu_doms, routed_doms);
    }
}