cotis-wgpu 0.1.0-alpha

//! Advanced extension API: UI geometry batch and WGSL render pipeline.
//!
//! Extension API for custom [`crate::drawable::WgpuDrawable`] implementations; not re-exported in
//! [`crate::prelude`] and may change between releases.
//!
//! [`GeometryBatch`] accumulates [`UIVertex`] data on the CPU and flushes it through a
//! [`UIPipeline`] (WGSL shader in `ui_shader.wgsl`) each frame. Vertex positions are in
//! physical window pixels with a z component for depth ordering.

use std::ops::{Add, Mul, Sub};

use bytemuck::{Pod, Zeroable};
use wgpu::util::DeviceExt;

/// Per-corner border radii for rounded rectangles, in physical pixels.
#[derive(Debug, Clone, Copy)]
pub struct UICornerRadii {
    /// Top-left corner radius.
    pub top_left: f32,
    /// Top-right corner radius.
    pub top_right: f32,
    /// Bottom-left corner radius.
    pub bottom_left: f32,
    /// Bottom-right corner radius.
    pub bottom_right: f32,
}

/// Normalized RGB color (0.0..=1.0) for UI geometry vertices.
#[derive(Copy, Clone, Debug, Pod, Zeroable)]
#[repr(C)]
pub struct UIColor {
    /// Red channel.
    pub r: f32,
    /// Green channel.
    pub g: f32,
    /// Blue channel.
    pub b: f32,
}

/// Per-edge border thickness for [`GeometryBatch::rectangle`], in physical pixels.
pub struct UIBorderThickness {
    /// Top edge thickness.
    pub top: f32,
    /// Left edge thickness.
    pub left: f32,
    /// Bottom edge thickness.
    pub bottom: f32,
    /// Right edge thickness.
    pub right: f32,
}

/// 3D position for UI vertices (x, y in pixels; z for depth ordering).
#[derive(Copy, Clone, Debug, Pod, Zeroable)]
#[repr(C)]
pub struct UIPosition {
    /// Horizontal position in physical pixels.
    pub x: f32,
    /// Vertical position in physical pixels.
    pub y: f32,
    /// Depth value for z-ordering (lower values draw on top with current pipeline settings).
    pub z: f32,
}

impl Default for UIPosition {
    fn default() -> Self {
        Self::new()
    }
}

impl UIPosition {
    /// Returns the origin position `(0, 0, 0)`.
    pub const fn new() -> Self {
        Self {
            x: 0.0,
            y: 0.0,
            z: 0.0,
        }
    }

    /// Rotates this position in-place by `degrees` (clockwise in screen space).
    pub fn rotate(&mut self, mut degrees: f32) {
        degrees = -degrees;
        degrees *= std::f32::consts::PI / 180.0;
        let (sn, cs) = degrees.sin_cos();
        *self = Self {
            x: self.x * cs - self.y * sn,
            y: self.x * sn + self.y * cs,
            z: self.z,
        };
    }

    /// Returns a copy with `x` offset by the given amount.
    pub fn with_x(self, x: f32) -> Self {
        Self {
            x: self.x + x,
            y: self.y,
            z: self.z,
        }
    }

    /// Returns a copy with `y` offset by the given amount.
    pub fn with_y(self, y: f32) -> Self {
        Self {
            x: self.x,
            y: self.y + y,
            z: self.z,
        }
    }
}

impl Add for UIPosition {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        Self {
            x: self.x + other.x,
            y: self.y + other.y,
            z: self.z,
        }
    }
}

impl Add<f32> for UIPosition {
    type Output = Self;

    fn add(self, rhs: f32) -> Self {
        Self {
            x: self.x + rhs,
            y: self.y + rhs,
            z: self.z,
        }
    }
}

impl Sub<f32> for UIPosition {
    type Output = Self;

    fn sub(self, rhs: f32) -> Self {
        Self {
            x: self.x - rhs,
            y: self.y - rhs,
            z: self.z,
        }
    }
}

impl Mul<f32> for UIPosition {
    type Output = Self;

    fn mul(self, rhs: f32) -> Self {
        Self {
            x: self.x * rhs,
            y: self.y * rhs,
            z: self.z,
        }
    }
}

/// Width and height pair stored in UI vertices for shader coordinate normalization.
#[derive(Copy, Clone, Debug, Pod, Zeroable)]
#[repr(C)]
pub struct UISize {
    /// Width in physical pixels.
    pub width: f32,
    /// Height in physical pixels.
    pub height: f32,
}

/// Single vertex for the UI geometry render pipeline.
#[derive(Copy, Clone, Debug, Pod, Zeroable)]
#[repr(C)]
pub struct UIVertex {
    /// Vertex position and depth.
    pub position: UIPosition,
    /// Vertex color (normalized RGB).
    pub color: UIColor,
    /// Viewport size used by the vertex shader for NDC conversion.
    pub size: UISize,
}

impl UIVertex {
    /// Creates a vertex with zero position/color and the given viewport size.
    pub fn new(size: (i32, i32)) -> Self {
        Self {
            position: UIPosition::new(),
            color: UIColor {
                r: 0.0,
                g: 0.0,
                b: 0.0,
            },
            size: UISize {
                width: size.0 as f32,
                height: size.1 as f32,
            },
        }
    }

    /// Returns the wgpu vertex buffer layout descriptor for [`UIVertex`].
    pub fn layout() -> wgpu::VertexBufferLayout<'static> {
        const ATTR: [wgpu::VertexAttribute; 3] =
            wgpu::vertex_attr_array![0 => Float32x3, 1 => Float32x3, 2 => Float32x2];

        wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<UIVertex>() as u64,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &ATTR,
        }
    }
}

/// CPU-side geometry batch flushed each frame through the UI render pipeline.
pub struct GeometryBatch {
    vertices: Vec<UIVertex>,
    buffer: wgpu::Buffer,
    number_of_vertices: usize,
    render_pipeline: wgpu::RenderPipeline,
}

impl GeometryBatch {
    /// Creates a geometry batch with pre-allocated vertex capacity and a UI render pipeline.
    pub fn new(device: &wgpu::Device, pixel_format: wgpu::TextureFormat, size: (i32, i32)) -> Self {
        let (buffer, vertices) = make_vertex_buffer(device, "cotis-wgpu ui vertices", 10_000, size);

        let mut pipeline_builder = UIPipeline::new(pixel_format);
        pipeline_builder.add_buffer_layout(UIVertex::layout());
        let render_pipeline = pipeline_builder.build_pipeline(device);

        Self {
            vertices,
            buffer,
            number_of_vertices: 0,
            render_pipeline,
        }
    }

    /// Updates the viewport size stored in all pre-allocated vertices.
    pub fn resize_vertices(&mut self, size: (i32, i32)) {
        for vertex in &mut self.vertices {
            vertex.size.width = size.0 as f32;
            vertex.size.height = size.1 as f32;
        }
    }

    /// Uploads accumulated vertices to the GPU and draws them.
    ///
    /// Resets the vertex count after drawing. No-op if no vertices were added.
    pub fn flush(&mut self, render_pass: &mut wgpu::RenderPass<'_>, queue: &wgpu::Queue) {
        if self.number_of_vertices == 0 {
            return;
        }

        render_pass.set_pipeline(&self.render_pipeline);
        queue.write_buffer(
            &self.buffer,
            0,
            bytemuck::cast_slice(&self.vertices[..self.number_of_vertices]),
        );
        render_pass.set_vertex_buffer(0, self.buffer.slice(..));
        render_pass.draw(0..self.number_of_vertices as u32, 0..1);
        self.number_of_vertices = 0;
    }

    /// Adds a filled triangle (3 vertices).
    pub fn triangle(&mut self, positions: &[UIPosition; 3], color: UIColor) {
        let Some(vertices) = self
            .vertices
            .get_mut(self.number_of_vertices..self.number_of_vertices + 3)
        else {
            return;
        };

        for (vertex, position) in vertices.iter_mut().zip(positions.iter()) {
            vertex.position = *position;
            vertex.color = color;
            self.number_of_vertices += 1;
        }
    }

    /// Adds a filled quad as two triangles (6 vertices).
    pub fn quad(&mut self, positions: &[UIPosition; 4], color: UIColor) {
        let Some(vertices) = self
            .vertices
            .get_mut(self.number_of_vertices..self.number_of_vertices + 6)
        else {
            return;
        };

        vertices[0].position = positions[0];
        vertices[0].color = color;
        vertices[1].position = positions[1];
        vertices[1].color = color;
        vertices[2].position = positions[2];
        vertices[2].color = color;
        vertices[3].position = positions[0];
        vertices[3].color = color;
        vertices[4].position = positions[2];
        vertices[4].color = color;
        vertices[5].position = positions[3];
        vertices[5].color = color;
        self.number_of_vertices += 6;
    }

    /// Adds a line segment as a thin quad rotated by `angle` degrees.
    pub fn line(
        &mut self,
        position: UIPosition,
        length: f32,
        angle: f32,
        thickness: f32,
        color: UIColor,
    ) {
        let mut line = [UIPosition::new(); 4];
        line[0].y -= thickness / 2.0;
        line[1].y += thickness / 2.0;
        line[2].x += length;
        line[2].y += thickness / 2.0;
        line[3].x += length;
        line[3].y -= thickness / 2.0;

        for point in &mut line {
            point.rotate(angle);
            *point = *point + position;
        }

        self.quad(&line, color);
    }

    /// Adds an arc outline approximated by line segments.
    pub fn arc(
        &mut self,
        origin: UIPosition,
        radius: f32,
        degree_begin: f32,
        degree_end: f32,
        thickness: f32,
        color: UIColor,
    ) {
        let arc_length = (degree_end - degree_begin).abs();
        let number_of_segments = 10.0;
        let arc_segment_length = arc_length / number_of_segments;
        let arc_segment_distance =
            (2.0 * std::f32::consts::PI * radius) * (arc_segment_length / 360.0);

        let mut arc_point = UIPosition::new();

        for i in 0..number_of_segments as i32 {
            arc_point.x = radius;
            arc_point.y = 0.0;
            arc_point.rotate(degree_begin + arc_segment_length * i as f32);
            arc_point = arc_point + origin;

            self.line(
                arc_point,
                arc_segment_distance,
                degree_begin + 90.0 + arc_segment_length * i as f32 + arc_segment_length / 2.0,
                thickness,
                color,
            );
        }
    }

    /// Adds a filled arc (pie slice) approximated by triangles.
    pub fn filled_arc(
        &mut self,
        origin: UIPosition,
        radius: f32,
        degree_begin: f32,
        degree_end: f32,
        color: UIColor,
    ) {
        let arc_length = (degree_end - degree_begin).abs();
        let number_of_segments = 10.0;
        let arc_segment_length = arc_length / number_of_segments;

        let mut current_point = UIPosition::new();
        let mut next_point = UIPosition::new();

        for i in 0..number_of_segments as i32 {
            current_point.x = radius;
            current_point.y = 0.0;
            current_point.rotate(degree_begin + arc_segment_length * (i as f32 + 1.0));

            next_point.x = radius;
            next_point.y = 0.0;
            next_point.rotate(degree_begin + arc_segment_length * i as f32);

            self.triangle(
                &[current_point + origin, origin, next_point + origin],
                color,
            );
        }
    }

    /// Adds a rounded rectangle border outline.
    pub fn rectangle(
        &mut self,
        position: UIPosition,
        size: UIPosition,
        thickness: UIBorderThickness,
        color: UIColor,
        radii: UICornerRadii,
    ) {
        self.arc(
            position + radii.top_left,
            radii.top_left,
            90.0,
            180.0,
            thickness.top,
            color,
        );
        self.arc(
            position
                .with_x(size.x - radii.top_right)
                .with_y(radii.top_right),
            radii.top_right,
            0.0,
            90.0,
            thickness.top,
            color,
        );
        self.arc(
            position
                .with_y(size.y - radii.bottom_left)
                .with_x(radii.bottom_left),
            radii.bottom_left,
            180.0,
            270.0,
            thickness.bottom,
            color,
        );
        self.arc(
            position + (size - radii.bottom_right),
            radii.bottom_right,
            270.0,
            360.0,
            thickness.bottom,
            color,
        );

        self.line(
            position.with_x(radii.top_left),
            size.x - (radii.top_left + radii.top_right),
            0.0,
            thickness.top,
            color,
        );
        self.line(
            position.with_y(radii.top_left),
            size.y - (radii.top_left + radii.bottom_left),
            270.0,
            thickness.left,
            color,
        );
        self.line(
            position.with_x(radii.bottom_left).with_y(size.y),
            size.x - (radii.bottom_left + radii.bottom_right),
            0.0,
            thickness.bottom,
            color,
        );
        self.line(
            position.with_x(size.x).with_y(radii.top_right),
            size.y - (radii.top_right + radii.bottom_right),
            270.0,
            thickness.right,
            color,
        );
    }

    /// Adds a filled rounded rectangle.
    pub fn filled_rectangle(
        &mut self,
        position: UIPosition,
        size: UIPosition,
        color: UIColor,
        radii: UICornerRadii,
    ) {
        self.filled_arc(
            position + radii.top_left,
            radii.top_left,
            90.0,
            180.0,
            color,
        );
        self.filled_arc(
            position
                .with_x(size.x - radii.top_right)
                .with_y(radii.top_right),
            radii.top_right,
            0.0,
            90.0,
            color,
        );
        self.filled_arc(
            position
                .with_y(size.y - radii.bottom_left)
                .with_x(radii.bottom_left),
            radii.bottom_left,
            180.0,
            270.0,
            color,
        );
        self.filled_arc(
            position + (size - radii.bottom_right),
            radii.bottom_right,
            270.0,
            360.0,
            color,
        );

        self.quad(
            &[
                position.with_x(radii.top_left),
                position + radii.top_left,
                position
                    .with_x(size.x - radii.top_right)
                    .with_y(radii.top_right),
                position.with_x(size.x - radii.top_right),
            ],
            color,
        );
        self.quad(
            &[
                position
                    .with_x(radii.bottom_left)
                    .with_y(size.y - radii.bottom_left),
                position.with_x(radii.bottom_left).with_y(size.y),
                position.with_x(size.x - radii.bottom_right).with_y(size.y),
                position
                    .with_x(size.x - radii.bottom_right)
                    .with_y(size.y - radii.bottom_right),
            ],
            color,
        );
        self.quad(
            &[
                position.with_y(radii.top_left),
                position.with_y(size.y - radii.bottom_left),
                position
                    .with_x(radii.bottom_left)
                    .with_y(size.y - radii.bottom_left),
                position + radii.top_left,
            ],
            color,
        );
        self.quad(
            &[
                position.with_x(size.x - radii.top_right),
                position
                    .with_x(size.x - radii.bottom_right)
                    .with_y(size.y - radii.bottom_right),
                position.with_x(size.x).with_y(size.y - radii.bottom_right),
                position.with_x(size.x).with_y(radii.top_right),
            ],
            color,
        );
        self.quad(
            &[
                position + radii.top_left,
                position
                    .with_x(radii.bottom_left)
                    .with_y(size.y - radii.bottom_left),
                position
                    .with_x(size.x - radii.bottom_right)
                    .with_y(size.y - radii.bottom_right),
                position
                    .with_x(size.x - radii.top_right)
                    .with_y(radii.top_right),
            ],
            color,
        );
    }
}

fn make_vertex_buffer(
    device: &wgpu::Device,
    label: &str,
    triangle_capacity: usize,
    size: (i32, i32),
) -> (wgpu::Buffer, Vec<UIVertex>) {
    let vertices = vec![UIVertex::new(size); triangle_capacity * 3];
    let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some(label),
        contents: bytemuck::cast_slice(&vertices),
        usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
    });
    (buffer, vertices)
}

/// Builder for the UI WGSL render pipeline.
pub struct UIPipeline {
    pixel_format: wgpu::TextureFormat,
    vertex_buffer_layouts: Vec<wgpu::VertexBufferLayout<'static>>,
}

impl UIPipeline {
    /// Creates a pipeline builder targeting the given swapchain pixel format.
    pub fn new(pixel_format: wgpu::TextureFormat) -> Self {
        Self {
            pixel_format,
            vertex_buffer_layouts: Vec::new(),
        }
    }

    /// Registers a vertex buffer layout (typically [`UIVertex::layout`]).
    pub fn add_buffer_layout(&mut self, layout: wgpu::VertexBufferLayout<'static>) {
        self.vertex_buffer_layouts.push(layout);
    }

    /// Compiles the WGSL shader and builds the wgpu render pipeline.
    pub fn build_pipeline(&self, device: &wgpu::Device) -> wgpu::RenderPipeline {
        let source_code = include_str!("ui_shader.wgsl");
        let shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("cotis-wgpu ui shader"),
            source: wgpu::ShaderSource::Wgsl(source_code.into()),
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("cotis-wgpu ui pipeline layout"),
            bind_group_layouts: &[],
            push_constant_ranges: &[],
        });

        let render_targets = [Some(wgpu::ColorTargetState {
            format: self.pixel_format,
            blend: Some(wgpu::BlendState::REPLACE),
            write_mask: wgpu::ColorWrites::ALL,
        })];

        device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("cotis-wgpu ui pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader_module,
                entry_point: Some("vs_main"),
                buffers: &self.vertex_buffer_layouts,
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                unclipped_depth: false,
                polygon_mode: wgpu::PolygonMode::Fill,
                conservative: false,
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader_module,
                entry_point: Some("fs_main"),
                targets: &render_targets,
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            }),
            depth_stencil: Some(wgpu::DepthStencilState {
                format: wgpu::TextureFormat::Depth32Float,
                depth_write_enabled: true,
                depth_compare: wgpu::CompareFunction::Always,
                stencil: wgpu::StencilState::default(),
                bias: wgpu::DepthBiasState::default(),
            }),
            multisample: wgpu::MultisampleState {
                count: 1,
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            multiview: None,
            cache: None,
        })
    }
}