dagx/tests/patterns/butterfly.rs

//! Butterfly pattern DAG tests

use crate::common::task_fn;
use dagx::{DagResult, DagRunner, TaskHandle};

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

#[tokio::test]
async fn test_butterfly_network() -> DagResult<()> {
    let dag = DagRunner::new();

    // Simplified butterfly network - avoid tuples entirely
    // Layer 0: inputs
    let inputs: Vec<_> = (0..8)
        .map(|i| {
            dag.add_task(task_fn::<(), _, _>(move |_: ()| i as f64))
                .into()
        })
        .collect();

    // Layer 1: butterfly operations (compute pairs separately)
    let mut layer1 = Vec::new();
    for i in 0..4 {
        let idx1 = i * 2;
        let idx2 = i * 2 + 1;
        // Sum operation
        layer1.push(
            dag.add_task(task_fn::<(f64, f64), _, _>(|(a, b): (&f64, &f64)| {
                (a + b) / 2.0
            }))
            .depends_on((&inputs[idx1], &inputs[idx2])),
        );
        // Difference operation
        layer1.push(
            dag.add_task(task_fn::<(f64, f64), _, _>(|(a, b): (&f64, &f64)| {
                (a - b) / 2.0
            }))
            .depends_on((&inputs[idx1], &inputs[idx2])),
        );
    }

    // Layer 2: crossover connections
    let layer2: Vec<_> = vec![
        dag.add_task(task_fn::<(f64, f64), _, _>(|(a, b): (&f64, &f64)| a + b))
            .depends_on((&layer1[0], &layer1[4])),
        dag.add_task(task_fn::<(f64, f64), _, _>(|(a, b): (&f64, &f64)| a - b))
            .depends_on((&layer1[1], &layer1[5])),
        dag.add_task(task_fn::<(f64, f64), _, _>(|(a, b): (&f64, &f64)| a + b))
            .depends_on((&layer1[2], &layer1[6])),
        dag.add_task(task_fn::<(f64, f64), _, _>(|(a, b): (&f64, &f64)| a - b))
            .depends_on((&layer1[3], &layer1[7])),
    ];

    dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
        .await?;

    // Verify some outputs exist
    for task in &layer2 {
        let _ = dag.get(task)?;
    }

    Ok(())
}

#[tokio::test]
async fn test_recursive_butterfly() -> DagResult<()> {
    let dag = DagRunner::new();

    fn build_butterfly(
        dag: &DagRunner,
        inputs: Vec<dagx::TaskHandle<i32>>,
    ) -> Vec<dagx::TaskHandle<i32>> {
        if inputs.len() <= 2 {
            return inputs;
        }

        let mid = inputs.len() / 2;
        let mut outputs = Vec::new();

        // Butterfly connections
        for i in 0..mid {
            let sum = dag
                .add_task(task_fn::<(i32, i32), _, _>(|(a, b): (&i32, &i32)| a + b))
                .depends_on((&inputs[i], &inputs[i + mid]));

            let diff = dag
                .add_task(task_fn::<(i32, i32), _, _>(|(a, b): (&i32, &i32)| {
                    (a - b).abs()
                }))
                .depends_on((&inputs[i], &inputs[i + mid]));

            outputs.push(sum);
            outputs.push(diff);
        }

        outputs
    }

    // Create 4 inputs
    let inputs: Vec<dagx::TaskHandle<i32>> = vec![10, 20, 30, 40]
        .into_iter()
        .map(|val| {
            let task = dag.add_task(task_fn::<(), _, _>(move |_: ()| val));
            task.into()
        })
        .collect();

    let outputs = build_butterfly(&dag, inputs);

    dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
        .await?;

    assert_eq!(outputs.len(), 4);
    for output in outputs {
        let _ = dag.get(output)?;
    }

    Ok(())
}

#[tokio::test]
async fn test_benes_network() -> DagResult<()> {
    // Simplified Benes network - avoid tuples
    let dag = DagRunner::new();

    // 4 inputs
    let inputs: Vec<_> = (0..4)
        .map(|i| {
            dag.add_task(task_fn::<(), _, _>(move |_: ()| 1 << i))
                .into()
        }) // Powers of 2
        .collect();

    // First stage: sort pairs
    let stage1: Vec<_> = vec![
        dag.add_task(task_fn::<(i32, i32), _, _>(|(&a, &b): (&i32, &i32)| {
            a.min(b)
        }))
        .depends_on((&inputs[0], &inputs[1])),
        dag.add_task(task_fn::<(i32, i32), _, _>(|(&a, &b): (&i32, &i32)| {
            a.max(b)
        }))
        .depends_on((&inputs[0], &inputs[1])),
        dag.add_task(task_fn::<(i32, i32), _, _>(|(&a, &b): (&i32, &i32)| {
            a.min(b)
        }))
        .depends_on((&inputs[2], &inputs[3])),
        dag.add_task(task_fn::<(i32, i32), _, _>(|(&a, &b): (&i32, &i32)| {
            a.max(b)
        }))
        .depends_on((&inputs[2], &inputs[3])),
    ];

    // Second stage with crossover
    let outputs: Vec<_> = vec![
        dag.add_task(task_fn::<(i32, i32), _, _>(|(a, b): (&i32, &i32)| a | b))
            .depends_on((&stage1[0], &stage1[2])),
        dag.add_task(task_fn::<(i32, i32), _, _>(|(a, b): (&i32, &i32)| a | b))
            .depends_on((&stage1[1], &stage1[3])),
    ];

    dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
        .await?;

    for output in outputs {
        let _ = dag.get(output)?;
    }

    Ok(())
}

#[tokio::test]
async fn test_shuffle_exchange_network() -> DagResult<()> {
    let dag = DagRunner::new();
    let counter = Arc::new(AtomicUsize::new(0));

    // Create 8 inputs
    let inputs: Vec<TaskHandle<_>> = (0..8)
        .map(|i| {
            let counter = counter.clone();
            dag.add_task(task_fn::<(), _, _>(move |_: ()| {
                let counter = counter.clone();
                counter.fetch_add(1, Ordering::Relaxed);
                i
            }))
            .into()
        })
        .collect();

    // Shuffle stage: bit rotation
    let shuffled: Vec<_> = (0..8)
        .map(|i| {
            // Shuffle pattern: rotate bits
            let source = ((i << 1) | (i >> 2)) & 0x7;
            dag.add_task(task_fn::<i32, _, _>(move |&x: &i32| x))
                .depends_on(inputs[source])
        })
        .collect();

    // Exchange stage: adjacent pairs
    let exchanged: Vec<_> = (0..4)
        .flat_map(|i| {
            let idx = i * 2;
            vec![
                dag.add_task(task_fn::<(i32, i32), _, _>(
                    |(&a, &b): (&i32, &i32)| if a > b { b } else { a },
                ))
                .depends_on((&shuffled[idx], &shuffled[idx + 1])),
                dag.add_task(task_fn::<(i32, i32), _, _>(
                    |(&a, &b): (&i32, &i32)| if a > b { a } else { b },
                ))
                .depends_on((&shuffled[idx], &shuffled[idx + 1])),
            ]
        })
        .collect();

    dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
        .await?;

    assert_eq!(counter.load(Ordering::Relaxed), 8);

    for task in exchanged {
        let _ = dag.get(task)?;
    }

    Ok(())
}