oxigdal-gpu 0.1.6

//! Integration tests for indirect compute dispatch.
//!
//! Pure-Rust tests (arithmetic helpers, struct layout) run unconditionally.
//! GPU-dependent tests wrap `GpuContext::new()` in `std::panic::catch_unwind`
//! and gracefully skip when no wgpu backend is available, matching the pattern
//! used in `texture_compress_test.rs` and `texture_resample_test.rs`.

#![allow(clippy::unwrap_used, missing_docs)]

use oxigdal_gpu::{
    DispatchIndirectArgs, IndirectDispatchBuffer, args_for_elements, dispatch_indirect_on_pass,
    workgroup_count_1d, workgroup_count_2d, workgroup_count_3d,
};

// ─────────────────────────────────────────────────────────────────────────────
// Helper: try to obtain a GPU context without panicking.
// ─────────────────────────────────────────────────────────────────────────────

fn try_gpu_context() -> Option<oxigdal_gpu::GpuContext> {
    use std::panic::AssertUnwindSafe;

    let result = std::panic::catch_unwind(AssertUnwindSafe(|| {
        pollster::block_on(oxigdal_gpu::GpuContext::new())
    }));

    match result {
        Ok(Ok(ctx)) => Some(ctx),
        _ => None,
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 1 — workgroup_count_1d: exact multiple → 1 workgroup
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_workgroup_count_1d_exact_multiple() {
    assert_eq!(
        workgroup_count_1d(64, 64),
        1,
        "64 elements / workgroup_size=64 must produce exactly 1 workgroup"
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 2 — workgroup_count_1d: remainder rounds up
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_workgroup_count_1d_remainder_rounds_up() {
    assert_eq!(
        workgroup_count_1d(65, 64),
        2,
        "65 elements / workgroup_size=64 must round up to 2 workgroups"
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 3 — workgroup_count_1d: zero elements → zero workgroups
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_workgroup_count_1d_zero_returns_zero() {
    assert_eq!(
        workgroup_count_1d(0, 64),
        0,
        "0 elements must produce 0 workgroups regardless of workgroup_size"
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 4 — workgroup_count_2d: (128, 64) with (16, 8) → (8, 8)
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_workgroup_count_2d_dimensions() {
    let (wg_x, wg_y) = workgroup_count_2d(128, 64, 16, 8);
    assert_eq!(wg_x, 8, "width=128 / wg_x=16 must yield 8 workgroups");
    assert_eq!(wg_y, 8, "height=64 / wg_y=8 must yield 8 workgroups");
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 5 — workgroup_count_3d: (10, 10, 10) with (4, 4, 4) → (3, 3, 3)
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_workgroup_count_3d_dimensions() {
    let (wg_x, wg_y, wg_z) = workgroup_count_3d((10, 10, 10), (4, 4, 4));
    assert_eq!(wg_x, 3, "10 / 4 must ceil to 3 workgroups (X)");
    assert_eq!(wg_y, 3, "10 / 4 must ceil to 3 workgroups (Y)");
    assert_eq!(wg_z, 3, "10 / 4 must ceil to 3 workgroups (Z)");
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 6 — args_for_elements: 128 elements / 64 → { x:2, y:1, z:1 }
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_args_for_elements_constructs_correctly() {
    let args = args_for_elements(128, 64);
    assert_eq!(
        args,
        DispatchIndirectArgs { x: 2, y: 1, z: 1 },
        "args_for_elements(128, 64) must produce {{x:2, y:1, z:1}}"
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 7 — DispatchIndirectArgs is repr(C) and exactly 12 bytes
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_dispatch_indirect_args_repr_c_size_12_bytes() {
    assert_eq!(
        std::mem::size_of::<DispatchIndirectArgs>(),
        12,
        "DispatchIndirectArgs must be exactly 12 bytes (repr(C), three u32 fields)"
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 8 — DispatchIndirectArgs::as_bytes produces correct little-endian layout
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_dispatch_indirect_args_repr_c_byte_layout() {
    let args = DispatchIndirectArgs { x: 1, y: 2, z: 3 };
    let bytes = args.as_bytes();
    let expected: [u8; 12] = [
        1, 0, 0, 0, // x = 1 (LE)
        2, 0, 0, 0, // y = 2 (LE)
        3, 0, 0, 0, // z = 3 (LE)
    ];
    assert_eq!(
        bytes, expected,
        "as_bytes() must produce little-endian [1,0,0,0, 2,0,0,0, 3,0,0,0]; \
         got {:?}",
        bytes
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 9 — IndirectDispatchBuffer::new when backend present
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_indirect_dispatch_buffer_new_when_backend_present() {
    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
        let ctx = match try_gpu_context() {
            Some(c) => c,
            None => return, // no GPU — skip
        };

        let args = DispatchIndirectArgs::new(4, 2, 1);
        let buf = match IndirectDispatchBuffer::new(&ctx, args) {
            Ok(b) => b,
            Err(e) => {
                eprintln!("IndirectDispatchBuffer::new failed (skip): {e}");
                return;
            }
        };

        assert_eq!(buf.capacity(), 1, "single-slot buffer must have capacity 1");
        assert_eq!(buf.offset(), 0, "byte offset for a fresh buffer must be 0");
        // buffer() should return a valid (non-null) reference; we verify this
        // indirectly by checking the wgpu size via the buffer descriptor.
        let _ = buf.buffer();
    }));
    let _ = result;
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 10 — IndirectDispatchBuffer::update when backend present
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_indirect_dispatch_buffer_update_when_backend_present() {
    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
        let ctx = match try_gpu_context() {
            Some(c) => c,
            None => return,
        };

        let initial_args = DispatchIndirectArgs::new(1, 1, 1);
        let buf = match IndirectDispatchBuffer::new(&ctx, initial_args) {
            Ok(b) => b,
            Err(e) => {
                eprintln!("IndirectDispatchBuffer::new failed (skip): {e}");
                return;
            }
        };

        // Updating slot 0 must not panic or return an error.
        let new_args = DispatchIndirectArgs::new(8, 4, 2);
        buf.update(&ctx, new_args);

        // Capacity and offset remain unchanged after update.
        assert_eq!(buf.capacity(), 1, "capacity must remain 1 after update");
        assert_eq!(buf.offset(), 0, "offset must remain 0 after update");
    }));
    let _ = result;
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 11 — IndirectDispatchBuffer::new_with_capacity(4) → capacity == 4
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_indirect_dispatch_buffer_with_capacity_when_backend_present() {
    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
        let ctx = match try_gpu_context() {
            Some(c) => c,
            None => return,
        };

        let buf = match IndirectDispatchBuffer::new_with_capacity(&ctx, 4) {
            Ok(b) => b,
            Err(e) => {
                eprintln!("new_with_capacity(4) failed (skip): {e}");
                return;
            }
        };

        assert_eq!(
            buf.capacity(),
            4,
            "buffer created with max_dispatches=4 must report capacity=4"
        );
        assert_eq!(buf.offset(), 0, "byte offset must be 0 for a fresh buffer");

        // update_at on a valid slot must succeed.
        let args = DispatchIndirectArgs::new(2, 2, 1);
        match buf.update_at(&ctx, 3, args) {
            Ok(()) => {}
            Err(e) => eprintln!("update_at(3) failed (skip): {e}"),
        }

        // update_at on an out-of-bounds slot must fail.
        let result = buf.update_at(&ctx, 4, args);
        assert!(
            result.is_err(),
            "update_at(4) on a capacity-4 buffer must return an error"
        );
    }));
    let _ = result;
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 12 — dispatch_indirect_on_pass executes a trivial compute shader
//           and writes 42 into an output buffer when a GPU backend is present.
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_dispatch_indirect_on_pass_executes_when_backend_present() {
    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
        let ctx = match try_gpu_context() {
            Some(c) => c,
            None => return,
        };

        // ── WGSL: write 42 into binding 0 ──────────────────────────────────
        let shader_src = r#"
@group(0) @binding(0) var<storage, read_write> out: u32;

@compute @workgroup_size(1)
fn main() {
    out = 42u;
}
"#;

        // Compile the shader module.
        let shader_module = ctx
            .device()
            .create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("indirect_dispatch_test_shader"),
                source: wgpu::ShaderSource::Wgsl(shader_src.into()),
            });

        // Build bind group layout: one read-write storage buffer at binding 0.
        let bgl = ctx
            .device()
            .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("indirect_dispatch_test_bgl"),
                entries: &[wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: false },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                }],
            });

        let pipeline_layout =
            ctx.device()
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                    label: Some("indirect_dispatch_test_layout"),
                    bind_group_layouts: &[Some(&bgl)],
                    immediate_size: 0,
                });

        let pipeline = ctx
            .device()
            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("indirect_dispatch_test_pipeline"),
                layout: Some(&pipeline_layout),
                module: &shader_module,
                entry_point: Some("main"),
                compilation_options: wgpu::PipelineCompilationOptions::default(),
                cache: None,
            });

        // Output buffer: 4 bytes (one u32), readable back to CPU via COPY_SRC.
        let output_buf = ctx.device().create_buffer(&wgpu::BufferDescriptor {
            label: Some("indirect_dispatch_test_output"),
            size: 4,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        // Staging buffer for readback.
        let staging_buf = ctx.device().create_buffer(&wgpu::BufferDescriptor {
            label: Some("indirect_dispatch_test_staging"),
            size: 4,
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });

        let bind_group = ctx.device().create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("indirect_dispatch_test_bg"),
            layout: &bgl,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: output_buf.as_entire_binding(),
            }],
        });

        // Indirect dispatch buffer: dispatch (1, 1, 1).
        let indirect_buf =
            match IndirectDispatchBuffer::new(&ctx, DispatchIndirectArgs::new(1, 1, 1)) {
                Ok(b) => b,
                Err(e) => {
                    eprintln!("IndirectDispatchBuffer::new failed (skip): {e}");
                    return;
                }
            };

        // Record the compute pass and copy to staging.
        let mut encoder = ctx
            .device()
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("indirect_dispatch_test_encoder"),
            });

        {
            let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("indirect_dispatch_test_pass"),
                timestamp_writes: None,
            });
            dispatch_indirect_on_pass(&mut compute_pass, &pipeline, &[&bind_group], &indirect_buf);
        }

        encoder.copy_buffer_to_buffer(&output_buf, 0, &staging_buf, 0, 4);
        ctx.queue().submit(std::iter::once(encoder.finish()));

        // Map the staging buffer and read back the result.
        let (tx, rx) = std::sync::mpsc::channel();
        staging_buf
            .slice(..)
            .map_async(wgpu::MapMode::Read, move |v| {
                let _ = tx.send(v);
            });

        // Poll until the map completes.
        let _poll_thread = ctx.spawn_poll_task();

        match rx.recv_timeout(std::time::Duration::from_secs(5)) {
            Ok(Ok(())) => {}
            Ok(Err(e)) => {
                eprintln!("staging buffer map failed (skip): {e}");
                return;
            }
            Err(e) => {
                eprintln!("timeout waiting for GPU readback (skip): {e}");
                return;
            }
        }

        let mapped = staging_buf.slice(..).get_mapped_range();
        let value = u32::from_le_bytes(
            mapped[0..4]
                .try_into()
                .expect("slice of 4 bytes is always valid"),
        );
        drop(mapped);
        staging_buf.unmap();

        assert_eq!(
            value, 42,
            "the compute shader must have written 42 into the output buffer; got {value}"
        );
    }));
    let _ = result;
}