henad-compute 0.3.0

//! The display half of a GPU model: a texture the model's compute shader writes, and the
//! fullscreen-triangle render pipeline that samples it into the viewport.
//!
//! The display knows nothing of the model. A model writes RGBA pixels into [`DisplayTarget::view`], and
//! gets a paintable [`GpuDisplay`].

use std::sync::Arc;

use crate::display_scale::display_dims;

/// The render-only half, passed to the UI inside a snapshot. The UI's paint callback holds an
/// `Arc` of this and does nothing but bind and draw, never touching simulation state.
///
/// The `Arc` makes teardown safe. An in-flight paint callback keeps the pipeline and
/// its texture alive even if the sim thread and its state are dropped mid-frame.
#[derive(Debug)]
pub struct GpuDisplay {
    /// Width of the underlying grid in cells, which the UI fits the aspect ratio to. The texture behind it can be
    /// smaller, as [`crate::display_scale`] caps it.
    pub width: u32,
    /// Height of the underlying grid in cells.
    pub height: u32,
    /// Pipeline that samples the texture into the viewport.
    pub render_pipeline: wgpu::RenderPipeline,
    /// Bind group of the texture and its sampler.
    pub render_bind_group: wgpu::BindGroup,
}

/// A display texture plus the [`GpuDisplay`] that samples it.
#[derive(Debug)]
pub struct DisplayTarget {
    /// Texture view, bound as a `texture_storage_2d<rgba8unorm, write>` in the model's display compute pass.
    pub view: wgpu::TextureView,
    /// Texture dimensions, which the model's display pass dispatches over.
    pub dims: (u32, u32),
    /// Render half, for the snapshot.
    pub display: Arc<GpuDisplay>,
}

/// Creates the display texture and the pipeline that samples it into `target_format`.
///
/// `width` and `height` are the grid dimensions in cells. The texture is capped by [`display_dims`].
#[cfg_attr(
    all(target_arch = "wasm32", target_feature = "atomics"),
    expect(
        clippy::arc_with_non_send_sync,
        reason = "a display holds wgpu handles, which atomics leave unsendable"
    )
)]
pub fn build_display_target(
    device: &wgpu::Device,
    target_format: wgpu::TextureFormat,
    width: u32,
    height: u32,
) -> DisplayTarget {
    let (tex_w, tex_h) = display_dims(width, height, device.limits().max_texture_dimension_2d);
    let texture = device.create_texture(&wgpu::TextureDescriptor {
        label: Some("henad_gpu_display_texture"),
        size: wgpu::Extent3d {
            width: tex_w,
            height: tex_h,
            depth_or_array_layers: 1,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D2,
        format: wgpu::TextureFormat::Rgba8Unorm,
        usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING,
        view_formats: &[],
    });
    let view = texture.create_view(&wgpu::TextureViewDescriptor::default());

    let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
        label: Some("henad_gpu_display_sampler"),
        mag_filter: wgpu::FilterMode::Nearest,
        min_filter: wgpu::FilterMode::Nearest,
        ..Default::default()
    });

    let shader = device.create_shader_module(wgpu::include_wgsl!("display.wgsl"));

    let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: Some("henad_gpu_display_bind_group_layout"),
        entries: &[
            wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Texture {
                    sample_type: wgpu::TextureSampleType::Float { filterable: true },
                    view_dimension: wgpu::TextureViewDimension::D2,
                    multisampled: false,
                },
                count: None,
            },
            wgpu::BindGroupLayoutEntry {
                binding: 1,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                count: None,
            },
        ],
    });

    let render_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("henad_gpu_display_bind_group"),
        layout: &bind_group_layout,
        entries: &[
            wgpu::BindGroupEntry {
                binding: 0,
                resource: wgpu::BindingResource::TextureView(&view),
            },
            wgpu::BindGroupEntry {
                binding: 1,
                resource: wgpu::BindingResource::Sampler(&sampler),
            },
        ],
    });

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

    let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
        label: Some("henad_gpu_display_pipeline"),
        layout: Some(&pipeline_layout),
        vertex: wgpu::VertexState {
            module: &shader,
            entry_point: Some("vs_main"),
            buffers: &[],
            compilation_options: wgpu::PipelineCompilationOptions::default(),
        },
        fragment: Some(wgpu::FragmentState {
            module: &shader,
            entry_point: Some("fs_main"),
            targets: &[Some(wgpu::ColorTargetState {
                format: target_format,
                blend: None,
                write_mask: wgpu::ColorWrites::ALL,
            })],
            compilation_options: wgpu::PipelineCompilationOptions::default(),
        }),
        primitive: wgpu::PrimitiveState::default(),
        depth_stencil: None,
        multisample: wgpu::MultisampleState::default(),
        multiview_mask: None,
        cache: None,
    });

    DisplayTarget {
        view,
        dims: (tex_w, tex_h),
        display: Arc::new(GpuDisplay {
            width,
            height,
            render_pipeline,
            render_bind_group,
        }),
    }
}