weirflow 0.1.0

use vyre::ir::Program;
use vyre_primitives::graph::csr_backward_traverse::csr_backward_traverse;
use vyre_primitives::graph::csr_forward_traverse::csr_forward_traverse;
use vyre_primitives::graph::program_graph::ProgramGraphShape;
use vyre_primitives::predicate::edge_kind;

/// Edge mask for CFG-only traversal.
pub const CFG_EDGE_MASK: u32 = edge_kind::CONTROL;

/// Build one forward CFG dataflow propagation step over control-flow edges.
#[must_use]
pub fn forward_cfg_step(
    shape: ProgramGraphShape,
    frontier_in: &str,
    frontier_out: &str,
) -> Program {
    csr_forward_traverse(shape, frontier_in, frontier_out, CFG_EDGE_MASK)
}

/// Build one backward CFG dataflow propagation step over control-flow edges.
#[must_use]
pub fn backward_cfg_step(
    shape: ProgramGraphShape,
    frontier_in: &str,
    frontier_out: &str,
) -> Program {
    csr_backward_traverse(shape, frontier_in, frontier_out, CFG_EDGE_MASK)
}

#[cfg(test)]
mod tests {
    use super::*;

    // ------------------------------------------------------------------
    // forward_cfg_step parity with csr_forward_traverse
    // ------------------------------------------------------------------

    #[test]
    fn forward_cfg_step_parity_empty_graph() {
        let shape = ProgramGraphShape::new(0, 0);
        let expected = csr_forward_traverse(shape, "in", "out", CFG_EDGE_MASK);
        let actual = forward_cfg_step(shape, "in", "out");
        assert_eq!(actual, expected);
    }

    #[test]
    fn forward_cfg_step_parity_single_node() {
        let shape = ProgramGraphShape::new(1, 0);
        let expected = csr_forward_traverse(shape, "frontier", "next", CFG_EDGE_MASK);
        let actual = forward_cfg_step(shape, "frontier", "next");
        assert_eq!(actual, expected);
    }

    #[test]
    fn forward_cfg_step_parity_two_nodes_one_edge() {
        let shape = ProgramGraphShape::new(2, 1);
        let expected = csr_forward_traverse(shape, "a", "b", CFG_EDGE_MASK);
        let actual = forward_cfg_step(shape, "a", "b");
        assert_eq!(actual, expected);
    }

    #[test]
    fn forward_cfg_step_parity_ten_nodes() {
        let shape = ProgramGraphShape::new(10, 15);
        let expected = csr_forward_traverse(shape, "in", "out", CFG_EDGE_MASK);
        let actual = forward_cfg_step(shape, "in", "out");
        assert_eq!(actual, expected);
    }

    #[test]
    fn forward_cfg_step_parity_large_graph() {
        let shape = ProgramGraphShape::new(1000, 5000);
        let expected = csr_forward_traverse(shape, "src", "dst", CFG_EDGE_MASK);
        let actual = forward_cfg_step(shape, "src", "dst");
        assert_eq!(actual, expected);
    }

    #[test]
    fn forward_cfg_step_parity_different_names() {
        let shape = ProgramGraphShape::new(5, 5);
        let names = [
            ("input", "output"),
            (" frontier_in ", " frontier_out "),
            ("α", "β"),
            ("x", "y"),
        ];
        for (i, o) in &names {
            let expected = csr_forward_traverse(shape, i, o, CFG_EDGE_MASK);
            let actual = forward_cfg_step(shape, i, o);
            assert_eq!(actual, expected);
        }
    }

    // ------------------------------------------------------------------
    // backward_cfg_step parity with csr_backward_traverse
    // ------------------------------------------------------------------

    #[test]
    fn backward_cfg_step_parity_empty_graph() {
        let shape = ProgramGraphShape::new(0, 0);
        let expected = csr_backward_traverse(shape, "in", "out", CFG_EDGE_MASK);
        let actual = backward_cfg_step(shape, "in", "out");
        assert_eq!(actual, expected);
    }

    #[test]
    fn backward_cfg_step_parity_single_node() {
        let shape = ProgramGraphShape::new(1, 0);
        let expected = csr_backward_traverse(shape, "frontier", "next", CFG_EDGE_MASK);
        let actual = backward_cfg_step(shape, "frontier", "next");
        assert_eq!(actual, expected);
    }

    #[test]
    fn backward_cfg_step_parity_two_nodes_one_edge() {
        let shape = ProgramGraphShape::new(2, 1);
        let expected = csr_backward_traverse(shape, "a", "b", CFG_EDGE_MASK);
        let actual = backward_cfg_step(shape, "a", "b");
        assert_eq!(actual, expected);
    }

    #[test]
    fn backward_cfg_step_parity_ten_nodes() {
        let shape = ProgramGraphShape::new(10, 15);
        let expected = csr_backward_traverse(shape, "in", "out", CFG_EDGE_MASK);
        let actual = backward_cfg_step(shape, "in", "out");
        assert_eq!(actual, expected);
    }

    #[test]
    fn backward_cfg_step_parity_large_graph() {
        let shape = ProgramGraphShape::new(1000, 5000);
        let expected = csr_backward_traverse(shape, "src", "dst", CFG_EDGE_MASK);
        let actual = backward_cfg_step(shape, "src", "dst");
        assert_eq!(actual, expected);
    }

    #[test]
    fn backward_cfg_step_parity_different_names() {
        let shape = ProgramGraphShape::new(5, 5);
        let names = [
            ("input", "output"),
            ("frontier_in", "frontier_out"),
            ("α", "β"),
            ("x", "y"),
        ];
        for (i, o) in &names {
            let expected = csr_backward_traverse(shape, i, o, CFG_EDGE_MASK);
            let actual = backward_cfg_step(shape, i, o);
            assert_eq!(actual, expected);
        }
    }

    // ------------------------------------------------------------------
    // Empty graph steps
    // ------------------------------------------------------------------

    #[test]
    fn forward_empty_graph_does_not_panic() {
        let shape = ProgramGraphShape::new(0, 0);
        let _ = forward_cfg_step(shape, "in", "out");
    }

    #[test]
    fn backward_empty_graph_does_not_panic() {
        let shape = ProgramGraphShape::new(0, 0);
        let _ = backward_cfg_step(shape, "in", "out");
    }

    #[test]
    fn forward_empty_graph_workgroup_size_is_sensible() {
        let shape = ProgramGraphShape::new(0, 0);
        let prog = forward_cfg_step(shape, "in", "out");
        // Workgroup size should be [1,1,1] or whatever the primitive sets.
        // We just assert it's non-zero in x to avoid degenerate programs.
        assert!(prog.workgroup_size[0] > 0);
    }

    #[test]
    fn backward_empty_graph_workgroup_size_is_sensible() {
        let shape = ProgramGraphShape::new(0, 0);
        let prog = backward_cfg_step(shape, "in", "out");
        assert!(prog.workgroup_size[0] > 0);
    }

    // ------------------------------------------------------------------
    // Single-node graph steps
    // ------------------------------------------------------------------

    #[test]
    fn forward_single_node_does_not_panic() {
        let shape = ProgramGraphShape::new(1, 0);
        let _ = forward_cfg_step(shape, "in", "out");
    }

    #[test]
    fn backward_single_node_does_not_panic() {
        let shape = ProgramGraphShape::new(1, 0);
        let _ = backward_cfg_step(shape, "in", "out");
    }

    #[test]
    fn forward_single_node_has_buffers() {
        let shape = ProgramGraphShape::new(1, 0);
        let prog = forward_cfg_step(shape, "in", "out");
        assert!(!prog.buffers.is_empty());
    }

    #[test]
    fn backward_single_node_has_buffers() {
        let shape = ProgramGraphShape::new(1, 0);
        let prog = backward_cfg_step(shape, "in", "out");
        assert!(!prog.buffers.is_empty());
    }

    // ------------------------------------------------------------------
    // Boundary checks (OOB node IDs / large shapes)
    // ------------------------------------------------------------------

    #[test]
    fn forward_u32_max_nodes_does_not_panic() {
        let shape = ProgramGraphShape::new(u32::MAX, 0);
        let _ = forward_cfg_step(shape, "in", "out");
    }

    #[test]
    fn backward_u32_max_nodes_does_not_panic() {
        let shape = ProgramGraphShape::new(u32::MAX, 0);
        let _ = backward_cfg_step(shape, "in", "out");
    }

    #[test]
    fn forward_max_edges_zero_nodes_does_not_panic() {
        let shape = ProgramGraphShape::new(0, u32::MAX);
        let _ = forward_cfg_step(shape, "in", "out");
    }

    #[test]
    fn backward_max_edges_zero_nodes_does_not_panic() {
        let shape = ProgramGraphShape::new(0, u32::MAX);
        let _ = backward_cfg_step(shape, "in", "out");
    }

    #[test]
    fn forward_large_both_does_not_panic() {
        let shape = ProgramGraphShape::new(1_000_000, 10_000_000);
        let _ = forward_cfg_step(shape, "in", "out");
    }

    #[test]
    fn backward_large_both_does_not_panic() {
        let shape = ProgramGraphShape::new(1_000_000, 10_000_000);
        let _ = backward_cfg_step(shape, "in", "out");
    }

    #[test]
    fn forward_one_node_many_edges() {
        let shape = ProgramGraphShape::new(1, 100);
        let expected = csr_forward_traverse(shape, "in", "out", CFG_EDGE_MASK);
        let actual = forward_cfg_step(shape, "in", "out");
        assert_eq!(actual, expected);
    }

    #[test]
    fn backward_one_node_many_edges() {
        let shape = ProgramGraphShape::new(1, 100);
        let expected = csr_backward_traverse(shape, "in", "out", CFG_EDGE_MASK);
        let actual = backward_cfg_step(shape, "in", "out");
        assert_eq!(actual, expected);
    }

    // ------------------------------------------------------------------
    // Structural properties
    // ------------------------------------------------------------------

    #[test]
    fn forward_and_backward_are_different() {
        let shape = ProgramGraphShape::new(10, 20);
        let fwd = forward_cfg_step(shape, "in", "out");
        let bwd = backward_cfg_step(shape, "in", "out");
        assert_ne!(fwd, bwd);
    }

    #[test]
    fn forward_same_inputs_produce_same_output() {
        let shape = ProgramGraphShape::new(5, 5);
        let a = forward_cfg_step(shape, "in", "out");
        let b = forward_cfg_step(shape, "in", "out");
        assert_eq!(a, b);
    }

    #[test]
    fn backward_same_inputs_produce_same_output() {
        let shape = ProgramGraphShape::new(5, 5);
        let a = backward_cfg_step(shape, "in", "out");
        let b = backward_cfg_step(shape, "in", "out");
        assert_eq!(a, b);
    }

    #[test]
    fn forward_different_shapes_produce_different_programs() {
        let a = forward_cfg_step(ProgramGraphShape::new(2, 1), "in", "out");
        let b = forward_cfg_step(ProgramGraphShape::new(3, 2), "in", "out");
        assert_ne!(a, b);
    }

    #[test]
    fn backward_different_shapes_produce_different_programs() {
        let a = backward_cfg_step(ProgramGraphShape::new(2, 1), "in", "out");
        let b = backward_cfg_step(ProgramGraphShape::new(3, 2), "in", "out");
        assert_ne!(a, b);
    }

    #[test]
    fn forward_different_frontier_names_produce_different_programs() {
        let shape = ProgramGraphShape::new(5, 5);
        let a = forward_cfg_step(shape, "in1", "out");
        let b = forward_cfg_step(shape, "in2", "out");
        assert_ne!(a, b);
    }

    #[test]
    fn backward_different_frontier_names_produce_different_programs() {
        let shape = ProgramGraphShape::new(5, 5);
        let a = backward_cfg_step(shape, "in", "out1");
        let b = backward_cfg_step(shape, "in", "out2");
        assert_ne!(a, b);
    }

    #[test]
    fn cfg_edge_mask_matches_control_const() {
        assert_eq!(CFG_EDGE_MASK, edge_kind::CONTROL);
    }

    #[test]
    fn forward_program_has_expected_buffer_names() {
        let shape = ProgramGraphShape::new(3, 3);
        let prog = forward_cfg_step(shape, "my_in", "my_out");
        let names: Vec<_> = prog.buffers.iter().map(|b| b.name.as_ref()).collect();
        // The program should reference the provided frontier names somewhere
        // in its buffers or entry nodes. At minimum the standard graph
        // buffers should be present.
        assert!(!names.is_empty());
    }

    #[test]
    fn backward_program_has_expected_buffer_names() {
        let shape = ProgramGraphShape::new(3, 3);
        let prog = backward_cfg_step(shape, "my_in", "my_out");
        let names: Vec<_> = prog.buffers.iter().map(|b| b.name.as_ref()).collect();
        assert!(!names.is_empty());
    }
}