Skip to main content

cotis_wgpu/
pipeline.rs

1//! Advanced extension API: UI geometry batch and WGSL render pipeline.
2//!
3//! Extension API for custom [`crate::drawable::WgpuDrawable`] implementations; not re-exported in
4//! [`crate::prelude`] and may change between releases.
5//!
6//! [`GeometryBatch`] accumulates [`UIVertex`] data on the CPU and flushes it through a
7//! [`UIPipeline`] (WGSL shader in `ui_shader.wgsl`) each frame. Vertex positions are in
8//! physical window pixels with a z component for depth ordering.
9
10use std::ops::{Add, Mul, Sub};
11
12use bytemuck::{Pod, Zeroable};
13use wgpu::util::DeviceExt;
14
15/// Per-corner border radii for rounded rectangles, in physical pixels.
16#[derive(Debug, Clone, Copy)]
17pub struct UICornerRadii {
18    /// Top-left corner radius.
19    pub top_left: f32,
20    /// Top-right corner radius.
21    pub top_right: f32,
22    /// Bottom-left corner radius.
23    pub bottom_left: f32,
24    /// Bottom-right corner radius.
25    pub bottom_right: f32,
26}
27
28/// Normalized RGB color (0.0..=1.0) for UI geometry vertices.
29#[derive(Copy, Clone, Debug, Pod, Zeroable)]
30#[repr(C)]
31pub struct UIColor {
32    /// Red channel.
33    pub r: f32,
34    /// Green channel.
35    pub g: f32,
36    /// Blue channel.
37    pub b: f32,
38}
39
40/// Per-edge border thickness for [`GeometryBatch::rectangle`], in physical pixels.
41pub struct UIBorderThickness {
42    /// Top edge thickness.
43    pub top: f32,
44    /// Left edge thickness.
45    pub left: f32,
46    /// Bottom edge thickness.
47    pub bottom: f32,
48    /// Right edge thickness.
49    pub right: f32,
50}
51
52/// 3D position for UI vertices (x, y in pixels; z for depth ordering).
53#[derive(Copy, Clone, Debug, Pod, Zeroable)]
54#[repr(C)]
55pub struct UIPosition {
56    /// Horizontal position in physical pixels.
57    pub x: f32,
58    /// Vertical position in physical pixels.
59    pub y: f32,
60    /// Depth value for z-ordering (lower values draw on top with current pipeline settings).
61    pub z: f32,
62}
63
64impl Default for UIPosition {
65    fn default() -> Self {
66        Self::new()
67    }
68}
69
70impl UIPosition {
71    /// Returns the origin position `(0, 0, 0)`.
72    pub const fn new() -> Self {
73        Self {
74            x: 0.0,
75            y: 0.0,
76            z: 0.0,
77        }
78    }
79
80    /// Rotates this position in-place by `degrees` (clockwise in screen space).
81    pub fn rotate(&mut self, mut degrees: f32) {
82        degrees = -degrees;
83        degrees *= std::f32::consts::PI / 180.0;
84        let (sn, cs) = degrees.sin_cos();
85        *self = Self {
86            x: self.x * cs - self.y * sn,
87            y: self.x * sn + self.y * cs,
88            z: self.z,
89        };
90    }
91
92    /// Returns a copy with `x` offset by the given amount.
93    pub fn with_x(self, x: f32) -> Self {
94        Self {
95            x: self.x + x,
96            y: self.y,
97            z: self.z,
98        }
99    }
100
101    /// Returns a copy with `y` offset by the given amount.
102    pub fn with_y(self, y: f32) -> Self {
103        Self {
104            x: self.x,
105            y: self.y + y,
106            z: self.z,
107        }
108    }
109}
110
111impl Add for UIPosition {
112    type Output = Self;
113
114    fn add(self, other: Self) -> Self {
115        Self {
116            x: self.x + other.x,
117            y: self.y + other.y,
118            z: self.z,
119        }
120    }
121}
122
123impl Add<f32> for UIPosition {
124    type Output = Self;
125
126    fn add(self, rhs: f32) -> Self {
127        Self {
128            x: self.x + rhs,
129            y: self.y + rhs,
130            z: self.z,
131        }
132    }
133}
134
135impl Sub<f32> for UIPosition {
136    type Output = Self;
137
138    fn sub(self, rhs: f32) -> Self {
139        Self {
140            x: self.x - rhs,
141            y: self.y - rhs,
142            z: self.z,
143        }
144    }
145}
146
147impl Mul<f32> for UIPosition {
148    type Output = Self;
149
150    fn mul(self, rhs: f32) -> Self {
151        Self {
152            x: self.x * rhs,
153            y: self.y * rhs,
154            z: self.z,
155        }
156    }
157}
158
159/// Width and height pair stored in UI vertices for shader coordinate normalization.
160#[derive(Copy, Clone, Debug, Pod, Zeroable)]
161#[repr(C)]
162pub struct UISize {
163    /// Width in physical pixels.
164    pub width: f32,
165    /// Height in physical pixels.
166    pub height: f32,
167}
168
169/// Single vertex for the UI geometry render pipeline.
170#[derive(Copy, Clone, Debug, Pod, Zeroable)]
171#[repr(C)]
172pub struct UIVertex {
173    /// Vertex position and depth.
174    pub position: UIPosition,
175    /// Vertex color (normalized RGB).
176    pub color: UIColor,
177    /// Viewport size used by the vertex shader for NDC conversion.
178    pub size: UISize,
179}
180
181impl UIVertex {
182    /// Creates a vertex with zero position/color and the given viewport size.
183    pub fn new(size: (i32, i32)) -> Self {
184        Self {
185            position: UIPosition::new(),
186            color: UIColor {
187                r: 0.0,
188                g: 0.0,
189                b: 0.0,
190            },
191            size: UISize {
192                width: size.0 as f32,
193                height: size.1 as f32,
194            },
195        }
196    }
197
198    /// Returns the wgpu vertex buffer layout descriptor for [`UIVertex`].
199    pub fn layout() -> wgpu::VertexBufferLayout<'static> {
200        const ATTR: [wgpu::VertexAttribute; 3] =
201            wgpu::vertex_attr_array![0 => Float32x3, 1 => Float32x3, 2 => Float32x2];
202
203        wgpu::VertexBufferLayout {
204            array_stride: std::mem::size_of::<UIVertex>() as u64,
205            step_mode: wgpu::VertexStepMode::Vertex,
206            attributes: &ATTR,
207        }
208    }
209}
210
211/// CPU-side geometry batch flushed each frame through the UI render pipeline.
212pub struct GeometryBatch {
213    vertices: Vec<UIVertex>,
214    buffer: wgpu::Buffer,
215    number_of_vertices: usize,
216    render_pipeline: wgpu::RenderPipeline,
217}
218
219impl GeometryBatch {
220    /// Creates a geometry batch with pre-allocated vertex capacity and a UI render pipeline.
221    pub fn new(device: &wgpu::Device, pixel_format: wgpu::TextureFormat, size: (i32, i32)) -> Self {
222        let (buffer, vertices) = make_vertex_buffer(device, "cotis-wgpu ui vertices", 10_000, size);
223
224        let mut pipeline_builder = UIPipeline::new(pixel_format);
225        pipeline_builder.add_buffer_layout(UIVertex::layout());
226        let render_pipeline = pipeline_builder.build_pipeline(device);
227
228        Self {
229            vertices,
230            buffer,
231            number_of_vertices: 0,
232            render_pipeline,
233        }
234    }
235
236    /// Updates the viewport size stored in all pre-allocated vertices.
237    pub fn resize_vertices(&mut self, size: (i32, i32)) {
238        for vertex in &mut self.vertices {
239            vertex.size.width = size.0 as f32;
240            vertex.size.height = size.1 as f32;
241        }
242    }
243
244    /// Uploads accumulated vertices to the GPU and draws them.
245    ///
246    /// Resets the vertex count after drawing. No-op if no vertices were added.
247    pub fn flush(&mut self, render_pass: &mut wgpu::RenderPass<'_>, queue: &wgpu::Queue) {
248        if self.number_of_vertices == 0 {
249            return;
250        }
251
252        render_pass.set_pipeline(&self.render_pipeline);
253        queue.write_buffer(
254            &self.buffer,
255            0,
256            bytemuck::cast_slice(&self.vertices[..self.number_of_vertices]),
257        );
258        render_pass.set_vertex_buffer(0, self.buffer.slice(..));
259        render_pass.draw(0..self.number_of_vertices as u32, 0..1);
260        self.number_of_vertices = 0;
261    }
262
263    /// Adds a filled triangle (3 vertices).
264    pub fn triangle(&mut self, positions: &[UIPosition; 3], color: UIColor) {
265        let Some(vertices) = self
266            .vertices
267            .get_mut(self.number_of_vertices..self.number_of_vertices + 3)
268        else {
269            return;
270        };
271
272        for (vertex, position) in vertices.iter_mut().zip(positions.iter()) {
273            vertex.position = *position;
274            vertex.color = color;
275            self.number_of_vertices += 1;
276        }
277    }
278
279    /// Adds a filled quad as two triangles (6 vertices).
280    pub fn quad(&mut self, positions: &[UIPosition; 4], color: UIColor) {
281        let Some(vertices) = self
282            .vertices
283            .get_mut(self.number_of_vertices..self.number_of_vertices + 6)
284        else {
285            return;
286        };
287
288        vertices[0].position = positions[0];
289        vertices[0].color = color;
290        vertices[1].position = positions[1];
291        vertices[1].color = color;
292        vertices[2].position = positions[2];
293        vertices[2].color = color;
294        vertices[3].position = positions[0];
295        vertices[3].color = color;
296        vertices[4].position = positions[2];
297        vertices[4].color = color;
298        vertices[5].position = positions[3];
299        vertices[5].color = color;
300        self.number_of_vertices += 6;
301    }
302
303    /// Adds a line segment as a thin quad rotated by `angle` degrees.
304    pub fn line(
305        &mut self,
306        position: UIPosition,
307        length: f32,
308        angle: f32,
309        thickness: f32,
310        color: UIColor,
311    ) {
312        let mut line = [UIPosition::new(); 4];
313        line[0].y -= thickness / 2.0;
314        line[1].y += thickness / 2.0;
315        line[2].x += length;
316        line[2].y += thickness / 2.0;
317        line[3].x += length;
318        line[3].y -= thickness / 2.0;
319
320        for point in &mut line {
321            point.rotate(angle);
322            *point = *point + position;
323        }
324
325        self.quad(&line, color);
326    }
327
328    /// Adds an arc outline approximated by line segments.
329    pub fn arc(
330        &mut self,
331        origin: UIPosition,
332        radius: f32,
333        degree_begin: f32,
334        degree_end: f32,
335        thickness: f32,
336        color: UIColor,
337    ) {
338        let arc_length = (degree_end - degree_begin).abs();
339        let number_of_segments = 10.0;
340        let arc_segment_length = arc_length / number_of_segments;
341        let arc_segment_distance =
342            (2.0 * std::f32::consts::PI * radius) * (arc_segment_length / 360.0);
343
344        let mut arc_point = UIPosition::new();
345
346        for i in 0..number_of_segments as i32 {
347            arc_point.x = radius;
348            arc_point.y = 0.0;
349            arc_point.rotate(degree_begin + arc_segment_length * i as f32);
350            arc_point = arc_point + origin;
351
352            self.line(
353                arc_point,
354                arc_segment_distance,
355                degree_begin + 90.0 + arc_segment_length * i as f32 + arc_segment_length / 2.0,
356                thickness,
357                color,
358            );
359        }
360    }
361
362    /// Adds a filled arc (pie slice) approximated by triangles.
363    pub fn filled_arc(
364        &mut self,
365        origin: UIPosition,
366        radius: f32,
367        degree_begin: f32,
368        degree_end: f32,
369        color: UIColor,
370    ) {
371        let arc_length = (degree_end - degree_begin).abs();
372        let number_of_segments = 10.0;
373        let arc_segment_length = arc_length / number_of_segments;
374
375        let mut current_point = UIPosition::new();
376        let mut next_point = UIPosition::new();
377
378        for i in 0..number_of_segments as i32 {
379            current_point.x = radius;
380            current_point.y = 0.0;
381            current_point.rotate(degree_begin + arc_segment_length * (i as f32 + 1.0));
382
383            next_point.x = radius;
384            next_point.y = 0.0;
385            next_point.rotate(degree_begin + arc_segment_length * i as f32);
386
387            self.triangle(
388                &[current_point + origin, origin, next_point + origin],
389                color,
390            );
391        }
392    }
393
394    /// Adds a rounded rectangle border outline.
395    pub fn rectangle(
396        &mut self,
397        position: UIPosition,
398        size: UIPosition,
399        thickness: UIBorderThickness,
400        color: UIColor,
401        radii: UICornerRadii,
402    ) {
403        self.arc(
404            position + radii.top_left,
405            radii.top_left,
406            90.0,
407            180.0,
408            thickness.top,
409            color,
410        );
411        self.arc(
412            position
413                .with_x(size.x - radii.top_right)
414                .with_y(radii.top_right),
415            radii.top_right,
416            0.0,
417            90.0,
418            thickness.top,
419            color,
420        );
421        self.arc(
422            position
423                .with_y(size.y - radii.bottom_left)
424                .with_x(radii.bottom_left),
425            radii.bottom_left,
426            180.0,
427            270.0,
428            thickness.bottom,
429            color,
430        );
431        self.arc(
432            position + (size - radii.bottom_right),
433            radii.bottom_right,
434            270.0,
435            360.0,
436            thickness.bottom,
437            color,
438        );
439
440        self.line(
441            position.with_x(radii.top_left),
442            size.x - (radii.top_left + radii.top_right),
443            0.0,
444            thickness.top,
445            color,
446        );
447        self.line(
448            position.with_y(radii.top_left),
449            size.y - (radii.top_left + radii.bottom_left),
450            270.0,
451            thickness.left,
452            color,
453        );
454        self.line(
455            position.with_x(radii.bottom_left).with_y(size.y),
456            size.x - (radii.bottom_left + radii.bottom_right),
457            0.0,
458            thickness.bottom,
459            color,
460        );
461        self.line(
462            position.with_x(size.x).with_y(radii.top_right),
463            size.y - (radii.top_right + radii.bottom_right),
464            270.0,
465            thickness.right,
466            color,
467        );
468    }
469
470    /// Adds a filled rounded rectangle.
471    pub fn filled_rectangle(
472        &mut self,
473        position: UIPosition,
474        size: UIPosition,
475        color: UIColor,
476        radii: UICornerRadii,
477    ) {
478        self.filled_arc(
479            position + radii.top_left,
480            radii.top_left,
481            90.0,
482            180.0,
483            color,
484        );
485        self.filled_arc(
486            position
487                .with_x(size.x - radii.top_right)
488                .with_y(radii.top_right),
489            radii.top_right,
490            0.0,
491            90.0,
492            color,
493        );
494        self.filled_arc(
495            position
496                .with_y(size.y - radii.bottom_left)
497                .with_x(radii.bottom_left),
498            radii.bottom_left,
499            180.0,
500            270.0,
501            color,
502        );
503        self.filled_arc(
504            position + (size - radii.bottom_right),
505            radii.bottom_right,
506            270.0,
507            360.0,
508            color,
509        );
510
511        self.quad(
512            &[
513                position.with_x(radii.top_left),
514                position + radii.top_left,
515                position
516                    .with_x(size.x - radii.top_right)
517                    .with_y(radii.top_right),
518                position.with_x(size.x - radii.top_right),
519            ],
520            color,
521        );
522        self.quad(
523            &[
524                position
525                    .with_x(radii.bottom_left)
526                    .with_y(size.y - radii.bottom_left),
527                position.with_x(radii.bottom_left).with_y(size.y),
528                position.with_x(size.x - radii.bottom_right).with_y(size.y),
529                position
530                    .with_x(size.x - radii.bottom_right)
531                    .with_y(size.y - radii.bottom_right),
532            ],
533            color,
534        );
535        self.quad(
536            &[
537                position.with_y(radii.top_left),
538                position.with_y(size.y - radii.bottom_left),
539                position
540                    .with_x(radii.bottom_left)
541                    .with_y(size.y - radii.bottom_left),
542                position + radii.top_left,
543            ],
544            color,
545        );
546        self.quad(
547            &[
548                position.with_x(size.x - radii.top_right),
549                position
550                    .with_x(size.x - radii.bottom_right)
551                    .with_y(size.y - radii.bottom_right),
552                position.with_x(size.x).with_y(size.y - radii.bottom_right),
553                position.with_x(size.x).with_y(radii.top_right),
554            ],
555            color,
556        );
557        self.quad(
558            &[
559                position + radii.top_left,
560                position
561                    .with_x(radii.bottom_left)
562                    .with_y(size.y - radii.bottom_left),
563                position
564                    .with_x(size.x - radii.bottom_right)
565                    .with_y(size.y - radii.bottom_right),
566                position
567                    .with_x(size.x - radii.top_right)
568                    .with_y(radii.top_right),
569            ],
570            color,
571        );
572    }
573}
574
575fn make_vertex_buffer(
576    device: &wgpu::Device,
577    label: &str,
578    triangle_capacity: usize,
579    size: (i32, i32),
580) -> (wgpu::Buffer, Vec<UIVertex>) {
581    let vertices = vec![UIVertex::new(size); triangle_capacity * 3];
582    let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
583        label: Some(label),
584        contents: bytemuck::cast_slice(&vertices),
585        usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
586    });
587    (buffer, vertices)
588}
589
590/// Builder for the UI WGSL render pipeline.
591pub struct UIPipeline {
592    pixel_format: wgpu::TextureFormat,
593    vertex_buffer_layouts: Vec<wgpu::VertexBufferLayout<'static>>,
594}
595
596impl UIPipeline {
597    /// Creates a pipeline builder targeting the given swapchain pixel format.
598    pub fn new(pixel_format: wgpu::TextureFormat) -> Self {
599        Self {
600            pixel_format,
601            vertex_buffer_layouts: Vec::new(),
602        }
603    }
604
605    /// Registers a vertex buffer layout (typically [`UIVertex::layout`]).
606    pub fn add_buffer_layout(&mut self, layout: wgpu::VertexBufferLayout<'static>) {
607        self.vertex_buffer_layouts.push(layout);
608    }
609
610    /// Compiles the WGSL shader and builds the wgpu render pipeline.
611    pub fn build_pipeline(&self, device: &wgpu::Device) -> wgpu::RenderPipeline {
612        let source_code = include_str!("ui_shader.wgsl");
613        let shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
614            label: Some("cotis-wgpu ui shader"),
615            source: wgpu::ShaderSource::Wgsl(source_code.into()),
616        });
617
618        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
619            label: Some("cotis-wgpu ui pipeline layout"),
620            bind_group_layouts: &[],
621            push_constant_ranges: &[],
622        });
623
624        let render_targets = [Some(wgpu::ColorTargetState {
625            format: self.pixel_format,
626            blend: Some(wgpu::BlendState::REPLACE),
627            write_mask: wgpu::ColorWrites::ALL,
628        })];
629
630        device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
631            label: Some("cotis-wgpu ui pipeline"),
632            layout: Some(&pipeline_layout),
633            vertex: wgpu::VertexState {
634                module: &shader_module,
635                entry_point: Some("vs_main"),
636                buffers: &self.vertex_buffer_layouts,
637                compilation_options: wgpu::PipelineCompilationOptions::default(),
638            },
639            primitive: wgpu::PrimitiveState {
640                topology: wgpu::PrimitiveTopology::TriangleList,
641                strip_index_format: None,
642                front_face: wgpu::FrontFace::Ccw,
643                cull_mode: Some(wgpu::Face::Back),
644                unclipped_depth: false,
645                polygon_mode: wgpu::PolygonMode::Fill,
646                conservative: false,
647            },
648            fragment: Some(wgpu::FragmentState {
649                module: &shader_module,
650                entry_point: Some("fs_main"),
651                targets: &render_targets,
652                compilation_options: wgpu::PipelineCompilationOptions::default(),
653            }),
654            depth_stencil: Some(wgpu::DepthStencilState {
655                format: wgpu::TextureFormat::Depth32Float,
656                depth_write_enabled: true,
657                depth_compare: wgpu::CompareFunction::Always,
658                stencil: wgpu::StencilState::default(),
659                bias: wgpu::DepthBiasState::default(),
660            }),
661            multisample: wgpu::MultisampleState {
662                count: 1,
663                mask: !0,
664                alpha_to_coverage_enabled: false,
665            },
666            multiview: None,
667            cache: None,
668        })
669    }
670}