codecraft 0.2.0

//! Draws what yakui painted, with wgpu. Stands in for `yakui-wgpu`, which is
//! pinned to an older wgpu than the rest of the renderer.
use std::collections::HashMap;

use bytemuck::{Pod, Zeroable};
use yakui::paint::TextureFormat as YakuiFormat;
use yakui::paint::{AddressMode, PaintDom, Pipeline, Texture, TextureChange, TextureFilter};
use yakui::{ManagedTextureId, TextureId, Yakui};

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
struct Vertex {
    position: [f32; 2],
    texcoord: [f32; 2],
    color: [f32; 4],
}

impl Vertex {
    const LAYOUT: wgpu::VertexBufferLayout<'static> = wgpu::VertexBufferLayout {
        array_stride: std::mem::size_of::<Self>() as wgpu::BufferAddress,
        step_mode: wgpu::VertexStepMode::Vertex,
        attributes: &wgpu::vertex_attr_array![0 => Float32x2, 1 => Float32x2, 2 => Float32x4],
    };
}

struct GpuTexture {
    texture: wgpu::Texture,
    bind_group: wgpu::BindGroup,
    size: (u32, u32),
    format: YakuiFormat,
}

struct Draw {
    indices: std::ops::Range<u32>,
    texture: Option<ManagedTextureId>,
    pipeline: Pipeline,
    clip: Option<yakui::geometry::Rect>,
}

/// Draws a [`Yakui`]'s paint output onto a frame.
pub struct UiRenderer {
    main_pipeline: wgpu::RenderPipeline,
    text_pipeline: wgpu::RenderPipeline,
    bind_group_layout: wgpu::BindGroupLayout,
    /// One white texel for untextured draws, so the shader's multiply is a no-op.
    blank: wgpu::BindGroup,
    textures: HashMap<ManagedTextureId, GpuTexture>,
    vertices: wgpu::Buffer,
    vertex_capacity: usize,
    indices: wgpu::Buffer,
    index_capacity: usize,
    draws: Vec<Draw>,
    vertex_data: Vec<Vertex>,
    index_data: Vec<u32>,
}

impl UiRenderer {
    pub fn new(device: &wgpu::Device, queue: &wgpu::Queue, format: wgpu::TextureFormat) -> Self {
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("ui shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("ui.wgsl").into()),
        });

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("ui 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 pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("ui pipeline layout"),
            bind_group_layouts: &[Some(&bind_group_layout)],
            immediate_size: 0,
        });

        let pipeline = |label: &str, fragment: &str| {
            device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
                label: Some(label),
                layout: Some(&pipeline_layout),
                vertex: wgpu::VertexState {
                    module: &shader,
                    entry_point: Some("vs_main"),
                    buffers: &[Some(Vertex::LAYOUT)],
                    compilation_options: wgpu::PipelineCompilationOptions::default(),
                },
                fragment: Some(wgpu::FragmentState {
                    module: &shader,
                    entry_point: Some(fragment),
                    targets: &[Some(wgpu::ColorTargetState {
                        format,
                        blend: Some(wgpu::BlendState::PREMULTIPLIED_ALPHA_BLENDING),
                        write_mask: wgpu::ColorWrites::ALL,
                    })],
                    compilation_options: wgpu::PipelineCompilationOptions::default(),
                }),
                primitive: wgpu::PrimitiveState {
                    topology: wgpu::PrimitiveTopology::TriangleList,
                    strip_index_format: None,
                    front_face: wgpu::FrontFace::Ccw,
                    cull_mode: None,
                    polygon_mode: wgpu::PolygonMode::Fill,
                    unclipped_depth: false,
                    conservative: false,
                },
                depth_stencil: None,
                multisample: wgpu::MultisampleState::default(),
                multiview_mask: None,
                cache: None,
            })
        };
        let main_pipeline = pipeline("ui main pipeline", "fs_main");
        let text_pipeline = pipeline("ui text pipeline", "fs_text");

        let blank = {
            let texture = device.create_texture(&wgpu::TextureDescriptor {
                label: Some("ui blank texture"),
                size: wgpu::Extent3d {
                    width: 1,
                    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::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
                view_formats: &[],
            });
            queue.write_texture(
                texture.as_image_copy(),
                &[255, 255, 255, 255],
                wgpu::TexelCopyBufferLayout {
                    offset: 0,
                    bytes_per_row: Some(4),
                    rows_per_image: Some(1),
                },
                wgpu::Extent3d {
                    width: 1,
                    height: 1,
                    depth_or_array_layers: 1,
                },
            );
            let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
            let sampler = device.create_sampler(&wgpu::SamplerDescriptor::default());
            device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("ui blank 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 vertex_capacity = 1024;
        let index_capacity = 2048;
        Self {
            main_pipeline,
            text_pipeline,
            bind_group_layout,
            blank,
            textures: HashMap::new(),
            vertices: vertex_buffer(device, vertex_capacity),
            vertex_capacity,
            indices: index_buffer(device, index_capacity),
            index_capacity,
            draws: Vec::new(),
            vertex_data: Vec::new(),
            index_data: Vec::new(),
        }
    }

    /// Draws this frame's yakui output over `view`, which already holds the rest of the frame.
    pub fn render(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        view: &wgpu::TextureView,
        yakui: &mut Yakui,
    ) {
        let paint = yakui.paint();
        self.update_textures(device, queue, paint);
        self.gather(paint);
        if self.draws.is_empty() {
            return;
        }

        self.upload(device, queue);
        let surface = paint.surface_size();
        let (surface_w, surface_h) = (surface.x.max(1.0) as u32, surface.y.max(1.0) as u32);

        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("ui render pass"),
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view,
                resolve_target: None,
                depth_slice: None,
                ops: wgpu::Operations {
                    // The 3D pass has already cleared and drawn the world.
                    load: wgpu::LoadOp::Load,
                    store: wgpu::StoreOp::Store,
                },
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
            multiview_mask: None,
        });
        pass.set_vertex_buffer(0, self.vertices.slice(..));
        pass.set_index_buffer(self.indices.slice(..), wgpu::IndexFormat::Uint32);

        let mut last_clip = None;
        let mut last_pipeline = None;
        for draw in &self.draws {
            if last_pipeline != Some(draw.pipeline) {
                last_pipeline = Some(draw.pipeline);
                match draw.pipeline {
                    Pipeline::Text => pass.set_pipeline(&self.text_pipeline),
                    _ => pass.set_pipeline(&self.main_pipeline),
                }
            }
            if draw.clip != last_clip {
                last_clip = draw.clip;
                match draw.clip {
                    Some(rect) => {
                        let x = rect.pos().x.max(0.0) as u32;
                        let y = rect.pos().y.max(0.0) as u32;
                        let right = (rect.max().x.max(0.0) as u32).min(surface_w);
                        let bottom = (rect.max().y.max(0.0) as u32).min(surface_h);
                        // wgpu rejects an empty scissor rect, so skip the draw instead.
                        if x >= right || y >= bottom {
                            last_clip = None;
                            continue;
                        }
                        pass.set_scissor_rect(x, y, right - x, bottom - y);
                    }
                    None => pass.set_scissor_rect(0, 0, surface_w, surface_h),
                }
            }
            let bind_group = draw
                .texture
                .and_then(|id| self.textures.get(&id))
                .map(|texture| &texture.bind_group)
                .unwrap_or(&self.blank);
            pass.set_bind_group(0, bind_group, &[]);
            pass.draw_indexed(draw.indices.clone(), 0, 0..1);
        }
    }

    fn update_textures(&mut self, device: &wgpu::Device, queue: &wgpu::Queue, paint: &PaintDom) {
        for (id, texture) in paint.textures() {
            if !self.textures.contains_key(&id) {
                let gpu = self.create_texture(device, queue, texture);
                self.textures.insert(id, gpu);
            }
        }
        for (id, change) in paint.texture_edits() {
            match change {
                TextureChange::Added => {
                    if let Some(texture) = paint.texture(id) {
                        let gpu = self.create_texture(device, queue, texture);
                        self.textures.insert(id, gpu);
                    }
                }
                TextureChange::Removed => {
                    self.textures.remove(&id);
                }
                TextureChange::Modified => {
                    let Some(texture) = paint.texture(id) else {
                        continue;
                    };
                    let same_shape = self.textures.get(&id).is_some_and(|gpu| {
                        gpu.size == (texture.size().x, texture.size().y)
                            && gpu.format == texture.format()
                    });
                    if same_shape {
                        write_texture(queue, &self.textures[&id].texture, texture);
                    } else {
                        let gpu = self.create_texture(device, queue, texture);
                        self.textures.insert(id, gpu);
                    }
                }
            }
        }
    }

    fn create_texture(
        &self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        texture: &Texture,
    ) -> GpuTexture {
        let size = texture.size();
        let gpu = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("ui texture"),
            size: wgpu::Extent3d {
                width: size.x,
                height: size.y,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu_format(texture.format()),
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        });
        write_texture(queue, &gpu, texture);

        let view = gpu.create_view(&wgpu::TextureViewDescriptor::default());
        let address = match texture.address_mode {
            AddressMode::ClampToEdge => wgpu::AddressMode::ClampToEdge,
            AddressMode::Repeat => wgpu::AddressMode::Repeat,
        };
        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("ui texture sampler"),
            address_mode_u: address,
            address_mode_v: address,
            mag_filter: wgpu_filter(texture.mag_filter),
            min_filter: wgpu_filter(texture.min_filter),
            ..Default::default()
        });
        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("ui texture bind group"),
            layout: &self.bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(&view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(&sampler),
                },
            ],
        });
        GpuTexture {
            texture: gpu,
            bind_group,
            size: (size.x, size.y),
            format: texture.format(),
        }
    }

    fn gather(&mut self, paint: &PaintDom) {
        self.vertex_data.clear();
        self.index_data.clear();
        self.draws.clear();

        for layer in paint.layers().iter() {
            for call in &layer.calls {
                if call.indices.is_empty() {
                    continue;
                }
                let base = self.vertex_data.len() as u32;
                let start = self.index_data.len() as u32;
                self.vertex_data
                    .extend(call.vertices.iter().map(|vertex| Vertex {
                        position: vertex.position.into(),
                        texcoord: vertex.texcoord.into(),
                        color: vertex.color.into(),
                    }));
                self.index_data
                    .extend(call.indices.iter().map(|&index| base + index as u32));
                let end = self.index_data.len() as u32;

                let texture = match call.texture {
                    Some(TextureId::Managed(id)) => Some(id),
                    _ => None,
                };
                self.draws.push(Draw {
                    indices: start..end,
                    texture,
                    pipeline: call.pipeline,
                    clip: call.clip,
                });
            }
        }
    }

    fn upload(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
        if self.vertex_data.len() > self.vertex_capacity {
            self.vertex_capacity = self.vertex_data.len().next_power_of_two();
            self.vertices = vertex_buffer(device, self.vertex_capacity);
        }
        if self.index_data.len() > self.index_capacity {
            self.index_capacity = self.index_data.len().next_power_of_two();
            self.indices = index_buffer(device, self.index_capacity);
        }
        queue.write_buffer(&self.vertices, 0, bytemuck::cast_slice(&self.vertex_data));
        queue.write_buffer(&self.indices, 0, bytemuck::cast_slice(&self.index_data));
    }
}

fn vertex_buffer(device: &wgpu::Device, capacity: usize) -> wgpu::Buffer {
    device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("ui vertices"),
        size: (capacity * std::mem::size_of::<Vertex>()) as u64,
        usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    })
}

fn index_buffer(device: &wgpu::Device, capacity: usize) -> wgpu::Buffer {
    device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("ui indices"),
        size: (capacity * std::mem::size_of::<u32>()) as u64,
        usage: wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    })
}

fn wgpu_format(format: YakuiFormat) -> wgpu::TextureFormat {
    match format {
        YakuiFormat::Rgba8Srgb => wgpu::TextureFormat::Rgba8UnormSrgb,
        YakuiFormat::R8 => wgpu::TextureFormat::R8Unorm,
        other => panic!("yakui asked for a texture format this renderer has no use for: {other:?}"),
    }
}

fn wgpu_filter(filter: TextureFilter) -> wgpu::FilterMode {
    match filter {
        TextureFilter::Linear => wgpu::FilterMode::Linear,
        TextureFilter::Nearest => wgpu::FilterMode::Nearest,
    }
}

// Premultiplies colour by alpha: the pipelines blend premultiplied, or soft edges go dark where they fade.
fn write_texture(queue: &wgpu::Queue, gpu: &wgpu::Texture, texture: &Texture) {
    let size = texture.size();
    let (bytes_per_pixel, data): (u32, std::borrow::Cow<'_, [u8]>) = match texture.format() {
        YakuiFormat::Rgba8Srgb => {
            let mut data = texture.data().to_vec();
            for pixel in data.as_chunks_mut::<4>().0 {
                let alpha = pixel[3] as u32;
                for channel in &mut pixel[..3] {
                    *channel = ((*channel as u32 * alpha + 255) >> 8) as u8;
                }
            }
            (4, data.into())
        }
        YakuiFormat::R8 => (1, texture.data().into()),
        other => panic!("yakui asked for a texture format this renderer has no use for: {other:?}"),
    };
    queue.write_texture(
        gpu.as_image_copy(),
        &data,
        wgpu::TexelCopyBufferLayout {
            offset: 0,
            bytes_per_row: Some(bytes_per_pixel * size.x),
            rows_per_image: Some(size.y),
        },
        wgpu::Extent3d {
            width: size.x,
            height: size.y,
            depth_or_array_layers: 1,
        },
    );
}