ai-app/iris/tests/color_space.rs

#![recursion_limit = "256"]

use iris::{harness::Harness, prelude::*};
use pollster::FutureExt;
use std::sync::OnceLock;
use wgpu::TextureFormat;

const SIZE: Vec2 = Vec2::new(2.0, 1.0);
const SOLID: Srgba8 = Srgba8::rgb(17, 127, 231);
const CHANGED_SOLID: Srgba8 = Srgba8::rgb(243, 139, 168);
const IMAGE: Srgba8 = Srgba8::rgb(205, 214, 244);

/// Covers the whole colour path rather than a conversion helper: an sRGB
/// literal enters the linear paint buffer, an sRGB image is sampled as
/// linear, the Iris shader returns both, and the sRGB attachment encodes the
/// stored bytes. Replacing only the paint-table entry also proves that a
/// theme change reaches an already-retained primitive.
#[test]
fn solid_paints_and_images_round_trip_through_an_srgb_target() {
    let gpu = Gpu::open();
    let mut harness = Harness::new(SIZE, 1.0);
    let solid = harness.rsc.ui.paints.add(SOLID);
    let bitmap = image::DynamicImage::ImageRgba8(image::RgbaImage::from_pixel(
        1,
        1,
        image::Rgba([IMAGE.r, IMAGE.g, IMAGE.b, IMAGE.a]),
    ));
    let bitmap = image::<iris::harness::HarnessRsc>(bitmap)(&mut harness.rsc);
    let root = (rect(solid.clone()).sized((1, 1)), bitmap)
        .span(Dir::RIGHT)
        .add_strong(&mut harness.rsc)
        .any();
    harness.state.set_root(&mut harness.rsc, root);
    harness.frame(0);

    let mut renderer =
        UiRenderNode::new(&gpu.device, &gpu.queue, TextureFormat::Rgba8UnormSrgb, SIZE)
            .expect("the Iris pipeline should accept an sRGB render target");
    let first = render(&gpu, &mut renderer, &mut harness);
    assert_pixel(first[0], SOLID, "linear paint buffer -> sRGB attachment");
    assert_pixel(first[1], IMAGE, "sRGB texture -> shader -> sRGB attachment");

    harness.rsc.ui.paints.set(&solid, CHANGED_SOLID);
    let changed = render(&gpu, &mut renderer, &mut harness);
    assert_pixel(
        changed[0],
        CHANGED_SOLID,
        "updated paint table -> retained primitive",
    );
    assert_pixel(changed[1], IMAGE, "unchanged image after paint update");
}

fn render(gpu: &Gpu, renderer: &mut UiRenderNode, harness: &mut Harness) -> [[u8; 4]; 2] {
    renderer.update(
        &gpu.device,
        &gpu.queue,
        &mut harness.rsc.ui,
        &mut harness.render,
    );
    let texture = gpu.device.create_texture(&wgpu::TextureDescriptor {
        label: Some("Iris colour-space target"),
        size: wgpu::Extent3d {
            width: 2,
            height: 1,
            depth_or_array_layers: 1,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D2,
        format: wgpu::TextureFormat::Rgba8UnormSrgb,
        usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
        view_formats: &[],
    });
    let view = texture.create_view(&Default::default());
    let readback = gpu.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("Iris colour-space readback"),
        size: 256,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    let mut encoder = gpu.device.create_command_encoder(&Default::default());
    {
        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("Iris colour-space pass"),
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: &view,
                depth_slice: None,
                resolve_target: None,
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Clear(LinearRgba::BLACK.to_wgpu()),
                    store: wgpu::StoreOp::Store,
                },
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
            multiview_mask: None,
        });
        renderer.draw(&mut pass);
    }
    encoder.copy_texture_to_buffer(
        texture.as_image_copy(),
        wgpu::TexelCopyBufferInfo {
            buffer: &readback,
            layout: wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(256),
                rows_per_image: Some(1),
            },
        },
        wgpu::Extent3d {
            width: 2,
            height: 1,
            depth_or_array_layers: 1,
        },
    );
    gpu.queue.submit([encoder.finish()]);

    let slice = readback.slice(..);
    slice.map_async(wgpu::MapMode::Read, |result| {
        result.expect("mapping the colour-space readback")
    });
    gpu.device
        .poll(wgpu::PollType::wait_indefinitely())
        .expect("waiting for the colour-space readback");
    let mapped = slice
        .get_mapped_range()
        .expect("reading the mapped colour-space buffer");
    let pixels = [
        mapped[0..4].try_into().unwrap(),
        mapped[4..8].try_into().unwrap(),
    ];
    drop(mapped);
    readback.unmap();
    pixels
}

fn assert_pixel(got: [u8; 4], want: Srgba8, path: &str) {
    let want = [want.r, want.g, want.b, want.a];
    assert!(
        got.into_iter()
            .zip(want)
            .all(|(got, want)| got.abs_diff(want) <= 1),
        "{path}: stored {got:?}, expected {want:?} (within one code value)",
    );
}

fn vulkan_instance() -> &'static wgpu::Instance {
    static INSTANCE: OnceLock<wgpu::Instance> = OnceLock::new();
    INSTANCE.get_or_init(wgpu::Instance::default)
}

struct Gpu {
    device: wgpu::Device,
    queue: wgpu::Queue,
}

impl Gpu {
    fn open() -> Self {
        let adapter = vulkan_instance()
            .request_adapter(&wgpu::RequestAdapterOptions::default())
            .block_on()
            .expect("no wgpu adapter, so Iris's colour-space pipeline went unchecked");
        let info = adapter.get_info();
        eprintln!(
            "color_space: {} ({:?}, {})",
            info.name, info.backend, info.driver
        );
        let (device, queue) = adapter
            .request_device(&wgpu::DeviceDescriptor {
                required_limits: iris_core::device_limits(),
                ..Default::default()
            })
            .block_on()
            .expect("could not get a device from the adapter");
        Self { device, queue }
    }
}