feat: migrate GPU renderer from WebGL to wgpu/WASM

1 files changed, 259 insertions, 0 deletions

diff --git a/rust/crates/gpu/src/mask_feather.rs b/rust/crates/gpu/src/mask_feather.rs
new file mode 100644
index 00000000..029b59ed
--- /dev/null
+++ b/rust/crates/gpu/src/mask_feather.rs
@@ -0,0 +1,259 @@
+use bytemuck::{Pod, Zeroable};
+use wgpu::util::DeviceExt;
+
+use crate::{GPU_TEXTURE_FORMAT, context::GpuContext};
+
+const FULLSCREEN_SHADER_SOURCE: &str = include_str!("shaders/fullscreen.wgsl");
+const JFA_DISTANCE_SHADER_SOURCE: &str = include_str!("shaders/jfa_distance.wgsl");
+
+pub struct ApplyMaskFeatherOptions<'a> {
+    pub mask: &'a wgpu::Texture,
+    pub width: u32,
+    pub height: u32,
+    pub feather: f32,
+}
+
+pub struct MaskFeatherPipeline {
+    inside_texture_bind_group_layout: wgpu::BindGroupLayout,
+    outside_texture_bind_group_layout: wgpu::BindGroupLayout,
+    uniform_bind_group_layout: wgpu::BindGroupLayout,
+    distance_pipeline: wgpu::RenderPipeline,
+}
+
+#[repr(C)]
+#[derive(Clone, Copy, Pod, Zeroable)]
+struct DistanceUniformBuffer {
+    resolution: [f32; 2],
+    feather_half: f32,
+    _padding: f32,
+}
+
+impl MaskFeatherPipeline {
+    pub fn new(device: &wgpu::Device) -> Self {
+        let inside_texture_bind_group_layout =
+            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+                label: Some("gpu-mask-distance-inside-layout"),
+                entries: &[
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 0,
+                        visibility: wgpu::ShaderStages::FRAGMENT,
+                        ty: wgpu::BindingType::Texture {
+                            multisampled: false,
+                            view_dimension: wgpu::TextureViewDimension::D2,
+                            sample_type: wgpu::TextureSampleType::Float { filterable: true },
+                        },
+                        count: None,
+                    },
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 1,
+                        visibility: wgpu::ShaderStages::FRAGMENT,
+                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
+                        count: None,
+                    },
+                ],
+            });
+        let outside_texture_bind_group_layout =
+            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+                label: Some("gpu-mask-distance-outside-layout"),
+                entries: &[
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 0,
+                        visibility: wgpu::ShaderStages::FRAGMENT,
+                        ty: wgpu::BindingType::Texture {
+                            multisampled: false,
+                            view_dimension: wgpu::TextureViewDimension::D2,
+                            sample_type: wgpu::TextureSampleType::Float { filterable: true },
+                        },
+                        count: None,
+                    },
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 1,
+                        visibility: wgpu::ShaderStages::FRAGMENT,
+                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
+                        count: None,
+                    },
+                ],
+            });
+        let uniform_bind_group_layout =
+            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+                label: Some("gpu-mask-distance-uniform-layout"),
+                entries: &[wgpu::BindGroupLayoutEntry {
+                    binding: 0,
+                    visibility: wgpu::ShaderStages::FRAGMENT,
+                    ty: wgpu::BindingType::Buffer {
+                        ty: wgpu::BufferBindingType::Uniform,
+                        has_dynamic_offset: false,
+                        min_binding_size: None,
+                    },
+                    count: None,
+                }],
+            });
+        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
+            label: Some("gpu-mask-distance-pipeline-layout"),
+            bind_group_layouts: &[
+                Some(&inside_texture_bind_group_layout),
+                Some(&outside_texture_bind_group_layout),
+                Some(&uniform_bind_group_layout),
+            ],
+            immediate_size: 0,
+        });
+        let vertex_shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
+            label: Some("gpu-mask-distance-fullscreen-shader"),
+            source: wgpu::ShaderSource::Wgsl(FULLSCREEN_SHADER_SOURCE.into()),
+        });
+        let fragment_shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
+            label: Some("gpu-mask-distance-fragment-shader"),
+            source: wgpu::ShaderSource::Wgsl(JFA_DISTANCE_SHADER_SOURCE.into()),
+        });
+        let distance_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
+            label: Some("gpu-mask-distance-pipeline"),
+            layout: Some(&pipeline_layout),
+            vertex: wgpu::VertexState {
+                module: &vertex_shader_module,
+                entry_point: Some("vertex_main"),
+                buffers: &[wgpu::VertexBufferLayout {
+                    array_stride: std::mem::size_of::<[f32; 2]>() as u64,
+                    step_mode: wgpu::VertexStepMode::Vertex,
+                    attributes: &[wgpu::VertexAttribute {
+                        format: wgpu::VertexFormat::Float32x2,
+                        offset: 0,
+                        shader_location: 0,
+                    }],
+                }],
+                compilation_options: wgpu::PipelineCompilationOptions::default(),
+            },
+            fragment: Some(wgpu::FragmentState {
+                module: &fragment_shader_module,
+                entry_point: Some("fragment_main"),
+                targets: &[Some(wgpu::ColorTargetState {
+                    format: GPU_TEXTURE_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,
+        });
+
+        Self {
+            inside_texture_bind_group_layout,
+            outside_texture_bind_group_layout,
+            uniform_bind_group_layout,
+            distance_pipeline,
+        }
+    }
+
+    pub fn apply_mask_feather(
+        &self,
+        context: &GpuContext,
+        ApplyMaskFeatherOptions {
+            mask,
+            width,
+            height,
+            feather,
+        }: ApplyMaskFeatherOptions<'_>,
+    ) -> wgpu::Texture {
+        let sdf = context
+            .sdf_pipeline()
+            .compute_signed_distance_field(context, mask, width, height);
+        let output_texture =
+            context.create_render_texture(width, height, "gpu-mask-feather-output");
+        let inside_view = sdf
+            .inside_texture
+            .create_view(&wgpu::TextureViewDescriptor::default());
+        let outside_view = sdf
+            .outside_texture
+            .create_view(&wgpu::TextureViewDescriptor::default());
+        let output_view = output_texture.create_view(&wgpu::TextureViewDescriptor::default());
+        let inside_bind_group = context
+            .device()
+            .create_bind_group(&wgpu::BindGroupDescriptor {
+                label: Some("gpu-mask-distance-inside-bind-group"),
+                layout: &self.inside_texture_bind_group_layout,
+                entries: &[
+                    wgpu::BindGroupEntry {
+                        binding: 0,
+                        resource: wgpu::BindingResource::TextureView(&inside_view),
+                    },
+                    wgpu::BindGroupEntry {
+                        binding: 1,
+                        resource: wgpu::BindingResource::Sampler(context.nearest_sampler()),
+                    },
+                ],
+            });
+        let outside_bind_group = context
+            .device()
+            .create_bind_group(&wgpu::BindGroupDescriptor {
+                label: Some("gpu-mask-distance-outside-bind-group"),
+                layout: &self.outside_texture_bind_group_layout,
+                entries: &[
+                    wgpu::BindGroupEntry {
+                        binding: 0,
+                        resource: wgpu::BindingResource::TextureView(&outside_view),
+                    },
+                    wgpu::BindGroupEntry {
+                        binding: 1,
+                        resource: wgpu::BindingResource::Sampler(context.nearest_sampler()),
+                    },
+                ],
+            });
+        let uniform_buffer = context
+            .device()
+            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
+                label: Some("gpu-mask-distance-uniform-buffer"),
+                contents: bytemuck::bytes_of(&DistanceUniformBuffer {
+                    resolution: [width as f32, height as f32],
+                    feather_half: feather / 2.0,
+                    _padding: 0.0,
+                }),
+                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
+            });
+        let uniform_bind_group = context
+            .device()
+            .create_bind_group(&wgpu::BindGroupDescriptor {
+                label: Some("gpu-mask-distance-uniform-bind-group"),
+                layout: &self.uniform_bind_group_layout,
+                entries: &[wgpu::BindGroupEntry {
+                    binding: 0,
+                    resource: uniform_buffer.as_entire_binding(),
+                }],
+            });
+        let mut encoder = context
+            .device()
+            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
+                label: Some("gpu-mask-distance-command-encoder"),
+            });
+
+        {
+            let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
+                label: Some("gpu-mask-distance-render-pass"),
+                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
+                    view: &output_view,
+                    resolve_target: None,
+                    depth_slice: None,
+                    ops: wgpu::Operations {
+                        load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
+                        store: wgpu::StoreOp::Store,
+                    },
+                })],
+                depth_stencil_attachment: None,
+                occlusion_query_set: None,
+                timestamp_writes: None,
+                multiview_mask: None,
+            });
+            render_pass.set_pipeline(&self.distance_pipeline);
+            render_pass.set_vertex_buffer(0, context.fullscreen_quad().slice(..));
+            render_pass.set_bind_group(0, &inside_bind_group, &[]);
+            render_pass.set_bind_group(1, &outside_bind_group, &[]);
+            render_pass.set_bind_group(2, &uniform_bind_group, &[]);
+            render_pass.draw(0..6, 0..1);
+        }
+
+        context.queue().submit([encoder.finish()]);
+        output_texture
+    }
+}