tessera_ui_basic_components/pipelines/
shape.rs

1mod command;
2use bytemuck::{Pod, Zeroable};
3use earcutr::earcut;
4use encase::ShaderType;
5use glam::Vec4;
6use log::error;
7use tessera_ui::{
8    PxPosition, PxSize,
9    renderer::DrawablePipeline,
10    wgpu::{self, include_wgsl, util::DeviceExt},
11};
12
13use crate::pipelines::pos_misc::pixel_to_ndc;
14
15use command::ShapeCommandComputed;
16
17pub use command::{RippleProps, ShadowProps, ShapeCommand};
18
19#[derive(ShaderType, Clone, Copy, Debug, PartialEq)]
20pub struct ShapeUniforms {
21    pub size_cr_border_width: Vec4,
22    pub primary_color: Vec4,
23    pub shadow_color: Vec4,
24    pub render_params: Vec4,
25    pub ripple_params: Vec4,
26    pub ripple_color: Vec4,
27    pub g2_k_value: f32,
28}
29
30/// Vertex for any shapes
31#[repr(C)]
32#[derive(Copy, Clone, Debug, Pod, Zeroable, PartialEq)]
33pub struct ShapeVertex {
34    /// Position of the vertex(x, y, z)
35    pub position: [f32; 3],
36    /// Color of the vertex
37    pub color: [f32; 3],
38    /// Normalized local position relative to rect center
39    pub local_pos: [f32; 2],
40}
41
42impl ShapeVertex {
43    /// Describe the vertex attributes
44    /// 0: position (x, y, z)
45    /// 1: color (r, g, b)
46    /// 2: local_pos (u, v)
47    /// The vertex attribute array is used to describe the vertex buffer layout
48    const ATTR: [wgpu::VertexAttribute; 3] =
49        wgpu::vertex_attr_array![0 => Float32x3, 1 => Float32x3, 2 => Float32x2];
50
51    /// Create a new vertex
52    fn new(pos: [f32; 2], color: [f32; 3], local_pos: [f32; 2]) -> Self {
53        Self {
54            position: [pos[0], pos[1], 0.0],
55            color,
56            local_pos,
57        }
58    }
59
60    /// Describe the vertex buffer layout
61    fn desc<'a>() -> wgpu::VertexBufferLayout<'a> {
62        wgpu::VertexBufferLayout {
63            array_stride: core::mem::size_of::<ShapeVertex>() as wgpu::BufferAddress,
64            step_mode: wgpu::VertexStepMode::Vertex,
65            attributes: &Self::ATTR,
66        }
67    }
68}
69
70pub struct ShapeVertexData<'a> {
71    pub polygon_vertices: &'a [[f32; 2]],
72    pub vertex_colors: &'a [[f32; 3]],
73    pub vertex_local_pos: &'a [[f32; 2]],
74}
75
76pub struct ShapePipeline {
77    pipeline: wgpu::RenderPipeline,
78    uniform_buffer: wgpu::Buffer,
79    #[allow(unused)]
80    bind_group_layout: wgpu::BindGroupLayout,
81    bind_group: wgpu::BindGroup,
82    shape_uniform_alignment: u32,
83    current_shape_uniform_offset: u32,
84    max_shape_uniform_buffer_offset: u32,
85}
86
87// Define MAX_CONCURRENT_SHAPES, can be adjusted later
88pub const MAX_CONCURRENT_SHAPES: wgpu::BufferAddress = 256;
89
90impl ShapePipeline {
91    pub fn new(gpu: &wgpu::Device, config: &wgpu::SurfaceConfiguration, sample_count: u32) -> Self {
92        let shader = gpu.create_shader_module(include_wgsl!("shape/shape.wgsl"));
93
94        let uniform_alignment =
95            gpu.limits().min_uniform_buffer_offset_alignment as wgpu::BufferAddress;
96        let size_of_shape_uniforms = std::mem::size_of::<ShapeUniforms>() as wgpu::BufferAddress;
97        let aligned_size_of_shape_uniforms =
98            wgpu::util::align_to(size_of_shape_uniforms, uniform_alignment);
99
100        let uniform_buffer = gpu.create_buffer(&wgpu::BufferDescriptor {
101            label: Some("Shape Uniform Buffer"),
102            size: MAX_CONCURRENT_SHAPES * aligned_size_of_shape_uniforms,
103            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
104            mapped_at_creation: false,
105        });
106
107        let bind_group_layout = gpu.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
108            entries: &[wgpu::BindGroupLayoutEntry {
109                binding: 0,
110                visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
111                ty: wgpu::BindingType::Buffer {
112                    ty: wgpu::BufferBindingType::Uniform,
113                    has_dynamic_offset: true, // Set to true for dynamic offsets
114                    min_binding_size: wgpu::BufferSize::new(
115                        std::mem::size_of::<ShapeUniforms>() as _
116                    ),
117                },
118                count: None,
119            }],
120            label: Some("shape_bind_group_layout"),
121        });
122
123        let bind_group = gpu.create_bind_group(&wgpu::BindGroupDescriptor {
124            layout: &bind_group_layout,
125            entries: &[wgpu::BindGroupEntry {
126                binding: 0,
127                resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
128                    buffer: &uniform_buffer,
129                    offset: 0, // Initial offset, will be overridden by dynamic offset
130                    size: wgpu::BufferSize::new(std::mem::size_of::<ShapeUniforms>() as _),
131                }),
132            }],
133            label: Some("shape_bind_group"),
134        });
135
136        let pipeline_layout = gpu.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
137            label: Some("Shape Pipeline Layout"),
138            bind_group_layouts: &[&bind_group_layout],
139            push_constant_ranges: &[],
140        });
141
142        let pipeline = gpu.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
143            label: Some("Shape Pipeline"),
144            layout: Some(&pipeline_layout),
145            vertex: wgpu::VertexState {
146                module: &shader,
147                entry_point: Some("vs_main"),
148                buffers: &[ShapeVertex::desc()],
149                compilation_options: Default::default(),
150            },
151            primitive: wgpu::PrimitiveState {
152                topology: wgpu::PrimitiveTopology::TriangleList,
153                strip_index_format: None,
154                front_face: wgpu::FrontFace::Ccw,
155                cull_mode: Some(wgpu::Face::Back),
156                unclipped_depth: false,
157                polygon_mode: wgpu::PolygonMode::Fill,
158                conservative: false,
159            },
160            depth_stencil: None,
161            multisample: wgpu::MultisampleState {
162                count: sample_count,
163                mask: !0,
164                alpha_to_coverage_enabled: false,
165            },
166            fragment: Some(wgpu::FragmentState {
167                module: &shader,
168                entry_point: Some("fs_main"),
169                compilation_options: Default::default(),
170                targets: &[Some(wgpu::ColorTargetState {
171                    format: config.format,
172                    blend: Some(wgpu::BlendState::ALPHA_BLENDING),
173                    write_mask: wgpu::ColorWrites::ALL,
174                })],
175            }),
176            multiview: None,
177            cache: None,
178        });
179
180        let size_of_shape_uniforms = std::mem::size_of::<ShapeUniforms>() as u32;
181        let alignment = gpu.limits().min_uniform_buffer_offset_alignment;
182        let shape_uniform_alignment =
183            wgpu::util::align_to(size_of_shape_uniforms, alignment) as u32;
184
185        let max_shape_uniform_buffer_offset =
186            (MAX_CONCURRENT_SHAPES as u32 - 1) * shape_uniform_alignment;
187
188        Self {
189            pipeline,
190            uniform_buffer,
191            bind_group_layout,
192            bind_group,
193            shape_uniform_alignment,
194            current_shape_uniform_offset: 0,
195            max_shape_uniform_buffer_offset,
196        }
197    }
198
199    fn draw_to_pass(
200        &self,
201        gpu: &wgpu::Device,
202        gpu_queue: &wgpu::Queue,
203        render_pass: &mut wgpu::RenderPass<'_>,
204        vertex_data_in: &ShapeVertexData,
205        uniforms: &ShapeUniforms,
206        dynamic_offset: u32,
207    ) {
208        let flat_polygon_vertices: Vec<f64> = vertex_data_in
209            .polygon_vertices
210            .iter()
211            .flat_map(|[x, y]| vec![*x as f64, *y as f64])
212            .collect();
213
214        let indices = earcut(&flat_polygon_vertices, &[], 2).unwrap_or_else(|e| {
215            error!("Earcut error: {e:?}");
216            Vec::new()
217        });
218
219        if indices.is_empty() && !vertex_data_in.polygon_vertices.is_empty() {
220            return;
221        }
222
223        let vertex_data: Vec<ShapeVertex> = indices
224            .iter()
225            .map(|&i| {
226                if i < vertex_data_in.polygon_vertices.len()
227                    && i < vertex_data_in.vertex_colors.len()
228                    && i < vertex_data_in.vertex_local_pos.len()
229                {
230                    ShapeVertex::new(
231                        vertex_data_in.polygon_vertices[i],
232                        vertex_data_in.vertex_colors[i],
233                        vertex_data_in.vertex_local_pos[i],
234                    )
235                } else {
236                    error!("Warning: Earcut index {i} out of bounds for input arrays.");
237                    // Fallback to the first vertex if index is out of bounds
238                    if !vertex_data_in.polygon_vertices.is_empty()
239                        && !vertex_data_in.vertex_colors.is_empty()
240                        && !vertex_data_in.vertex_local_pos.is_empty()
241                    {
242                        ShapeVertex::new(
243                            vertex_data_in.polygon_vertices[0],
244                            vertex_data_in.vertex_colors[0],
245                            vertex_data_in.vertex_local_pos[0],
246                        )
247                    } else {
248                        // This case should ideally not happen if inputs are validated
249                        // Or handle it by returning early / logging a more severe error
250                        ShapeVertex::new([0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0])
251                        // Placeholder
252                    }
253                }
254            })
255            .collect();
256
257        if vertex_data.is_empty() {
258            return;
259        }
260
261        let vertex_buffer = gpu.create_buffer_init(&wgpu::util::BufferInitDescriptor {
262            label: Some("Triangulated Vertex Buffer"),
263            contents: bytemuck::cast_slice(&vertex_data),
264            usage: wgpu::BufferUsages::VERTEX,
265        });
266
267        let mut buffer = encase::UniformBuffer::new(Vec::<u8>::new());
268        buffer.write(uniforms).unwrap();
269        let inner = buffer.into_inner();
270        gpu_queue.write_buffer(
271            &self.uniform_buffer,
272            dynamic_offset as wgpu::BufferAddress,
273            &inner,
274        );
275
276        render_pass.set_pipeline(&self.pipeline);
277        render_pass.set_bind_group(0, &self.bind_group, &[dynamic_offset]);
278        render_pass.set_vertex_buffer(0, vertex_buffer.slice(..));
279        render_pass.draw(0..vertex_data.len() as u32, 0..1);
280    }
281}
282
283#[allow(unused_variables)]
284impl DrawablePipeline<ShapeCommand> for ShapePipeline {
285    fn begin_frame(
286        &mut self,
287        _gpu: &wgpu::Device,
288        _gpu_queue: &wgpu::Queue,
289        _config: &wgpu::SurfaceConfiguration,
290    ) {
291        self.current_shape_uniform_offset = 0;
292    }
293
294    fn draw(
295        &mut self,
296        gpu: &wgpu::Device,
297        gpu_queue: &wgpu::Queue,
298        config: &wgpu::SurfaceConfiguration,
299        render_pass: &mut wgpu::RenderPass<'_>,
300        command: &ShapeCommand,
301        size: PxSize,
302        start_pos: PxPosition,
303        _scene_texture_view: &wgpu::TextureView,
304    ) {
305        // --- Fallback for ALL shapes, or primary path for non-G2 shapes ---
306        let computed_command = ShapeCommandComputed::from_command(command.clone(), size, start_pos);
307        let positions: Vec<[f32; 2]> = computed_command
308            .vertices
309            .iter()
310            .map(|v| {
311                pixel_to_ndc(
312                    PxPosition::from_f32_arr3(v.position),
313                    [config.width, config.height],
314                )
315            })
316            .collect();
317        let colors: Vec<[f32; 3]> = computed_command.vertices.iter().map(|v| v.color).collect();
318        let local_positions: Vec<[f32; 2]> = computed_command
319            .vertices
320            .iter()
321            .map(|v| v.local_pos)
322            .collect();
323
324        // Check if shadow needs to be drawn
325        let has_shadow = computed_command.uniforms.shadow_color[3] > 0.0
326            && computed_command.uniforms.render_params[2] > 0.0;
327
328        if has_shadow {
329            let dynamic_offset = self.current_shape_uniform_offset;
330            if dynamic_offset > self.max_shape_uniform_buffer_offset {
331                panic!(
332                    "Shape uniform buffer overflow for shadow: offset {} > max {}",
333                    dynamic_offset, self.max_shape_uniform_buffer_offset
334                );
335            }
336
337            let mut uniforms_for_shadow = computed_command.uniforms;
338            uniforms_for_shadow.render_params[3] = 2.0;
339
340            let vertex_data_for_shadow = ShapeVertexData {
341                polygon_vertices: &positions,
342                vertex_colors: &colors,
343                vertex_local_pos: &local_positions,
344            };
345
346            self.draw_to_pass(
347                gpu,
348                gpu_queue,
349                render_pass,
350                &vertex_data_for_shadow,
351                &uniforms_for_shadow,
352                dynamic_offset,
353            );
354            self.current_shape_uniform_offset += self.shape_uniform_alignment;
355        }
356
357        let dynamic_offset = self.current_shape_uniform_offset;
358        if dynamic_offset > self.max_shape_uniform_buffer_offset {
359            panic!(
360                "Shape uniform buffer overflow for object: offset {} > max {}",
361                dynamic_offset, self.max_shape_uniform_buffer_offset
362            );
363        }
364
365        let vertex_data_for_object = ShapeVertexData {
366            polygon_vertices: &positions,
367            vertex_colors: &colors,
368            vertex_local_pos: &local_positions,
369        };
370
371        self.draw_to_pass(
372            gpu,
373            gpu_queue,
374            render_pass,
375            &vertex_data_for_object,
376            &computed_command.uniforms,
377            dynamic_offset,
378        );
379        self.current_shape_uniform_offset += self.shape_uniform_alignment;
380    }
381}