Skip to main content

ranim_render/primitives/
vitems.rs

1use crate::utils::{WgpuContext, WgpuVecBuffer};
2use bytemuck::{Pod, Zeroable};
3use glam::{Vec3, Vec4};
4use ranim_core::{
5    components::{rgba::Rgba, width::Width},
6    core_item::vitem::{VItem, vitem_normal_from_points},
7};
8
9/// Per-item metadata stored in a GPU buffer.
10/// Tells shaders where each VItem's data lives in the merged buffers.
11#[repr(C)]
12#[derive(Debug, Default, Clone, Copy, Pod, Zeroable)]
13pub struct ItemInfo {
14    /// Offset into the merged points buffer
15    pub point_offset: u32,
16    /// Number of points for this item
17    pub point_count: u32,
18    /// Offset into the merged attribute buffers (fill_rgbas, stroke_rgbas, stroke_widths)
19    pub attr_offset: u32,
20    /// Number of attributes (= point_count.div_ceil(2))
21    pub attr_count: u32,
22}
23
24/// Per-item plane data (normal + origin), stored as array of structs.
25/// The origin is the first point of the item (used by vertex shader).
26/// basis_u/basis_v are generated deterministically from the normal in the shader.
27#[repr(C)]
28#[derive(Debug, Default, Clone, Copy, Pod, Zeroable)]
29pub struct PlaneData {
30    pub normal: Vec4, // xyz = normal, w = pad
31    pub origin: Vec4, // xyz = first point, w = pad
32}
33
34/// Merged GPU buffers for all VItems in a frame.
35///
36/// Instead of one set of buffers per VItem, all data is packed into
37/// contiguous arrays with an index table (`item_infos`) that tells
38/// shaders where each item's data lives.
39#[derive(bevy_ecs::prelude::Resource)]
40pub struct VItemsBuffer {
41    /// Per-item metadata: offsets and counts
42    pub(crate) item_infos_buffer: WgpuVecBuffer<ItemInfo>,
43    /// Per-item plane data (normal + origin for vertex shader)
44    pub(crate) planes_buffer: WgpuVecBuffer<PlaneData>,
45    /// Per-item clip boxes (5 i32 each: min_x, max_x, min_y, max_y, max_w)
46    pub(crate) clip_boxes_buffer: WgpuVecBuffer<i32>,
47
48    /// Merged 3D points from all VItems
49    pub(crate) points3d_buffer: WgpuVecBuffer<Vec4>,
50    /// Merged 2D projected points (written by compute shader)
51    pub(crate) points2d_buffer: WgpuVecBuffer<Vec4>,
52    /// Merged fill colors
53    pub(crate) fill_rgbas_buffer: WgpuVecBuffer<Rgba>,
54    /// Merged stroke colors
55    pub(crate) stroke_rgbas_buffer: WgpuVecBuffer<Rgba>,
56    /// Merged stroke widths
57    pub(crate) stroke_widths_buffer: WgpuVecBuffer<Width>,
58
59    /// Number of items
60    pub(crate) item_count: u32,
61    /// Total number of points across all items
62    pub(crate) total_points: u32,
63
64    /// Compute bind group (recreated when buffers resize)
65    pub(crate) compute_bind_group: Option<wgpu::BindGroup>,
66    /// Render bind group (recreated when buffers resize)
67    pub(crate) render_bind_group: Option<wgpu::BindGroup>,
68}
69
70impl VItemsBuffer {
71    pub fn new(ctx: &WgpuContext) -> Self {
72        // Start with empty buffers (minimum size 1 to avoid zero-size buffer)
73        let storage_rw = wgpu::BufferUsages::STORAGE
74            | wgpu::BufferUsages::COPY_DST
75            | wgpu::BufferUsages::COPY_SRC;
76        let storage_ro = wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST;
77
78        Self {
79            item_infos_buffer: WgpuVecBuffer::new(ctx, Some("Merged ItemInfos"), storage_ro, 1),
80            planes_buffer: WgpuVecBuffer::new(ctx, Some("Merged Planes"), storage_ro, 1),
81            clip_boxes_buffer: WgpuVecBuffer::new(ctx, Some("Merged ClipBoxes"), storage_rw, 5),
82            points3d_buffer: WgpuVecBuffer::new(ctx, Some("Merged Points3D"), storage_ro, 1),
83            points2d_buffer: WgpuVecBuffer::new(ctx, Some("Merged Points2D"), storage_rw, 1),
84            fill_rgbas_buffer: WgpuVecBuffer::new(ctx, Some("Merged FillRgbas"), storage_ro, 1),
85            stroke_rgbas_buffer: WgpuVecBuffer::new(ctx, Some("Merged StrokeRgbas"), storage_ro, 1),
86            stroke_widths_buffer: WgpuVecBuffer::new(
87                ctx,
88                Some("Merged StrokeWidths"),
89                storage_ro,
90                1,
91            ),
92            item_count: 0,
93            total_points: 0,
94            compute_bind_group: None,
95            render_bind_group: None,
96        }
97    }
98
99    /// Pack all VItems into the merged buffers. Called once per frame.
100    pub fn update<'a, I>(&mut self, ctx: &WgpuContext, vitems: I)
101    where
102        I: IntoIterator<Item = &'a VItem>,
103        I::IntoIter: ExactSizeIterator + Clone,
104    {
105        let vitems = vitems.into_iter();
106        if vitems.len() == 0 {
107            self.item_count = 0;
108            self.total_points = 0;
109            return;
110        }
111
112        let item_count = vitems.len();
113
114        // Pre-calculate total sizes
115        let total_points: usize = vitems.clone().map(|v| v.points.len()).sum();
116        let total_attrs: usize = vitems.clone().map(|v| v.points.len().div_ceil(2)).sum();
117
118        // Build index table and collect data
119        let mut item_infos = Vec::with_capacity(item_count);
120        let mut planes = Vec::with_capacity(item_count);
121        let mut all_points3d = Vec::with_capacity(total_points);
122        let mut all_fill_rgbas = Vec::with_capacity(total_attrs);
123        let mut all_stroke_rgbas = Vec::with_capacity(total_attrs);
124        let mut all_stroke_widths = Vec::with_capacity(total_attrs);
125
126        let mut point_offset: u32 = 0;
127        let mut attr_offset: u32 = 0;
128
129        for vitem in vitems {
130            let pc = vitem.points.len() as u32;
131            let ac = pc.div_ceil(2);
132
133            item_infos.push(ItemInfo {
134                point_offset,
135                point_count: pc,
136                attr_offset,
137                attr_count: ac,
138            });
139
140            let normal = vitem
141                .normal
142                .unwrap_or_else(|| vitem_normal_from_points(&vitem.points));
143            let origin = Vec3::new(vitem.points[0].x, vitem.points[0].y, vitem.points[0].z);
144            planes.push(PlaneData {
145                normal: Vec4::from((normal, 0.0)),
146                origin: Vec4::from((origin, 0.0)),
147            });
148
149            all_points3d.extend_from_slice(&vitem.points);
150            all_fill_rgbas.extend_from_slice(&vitem.fill_rgbas);
151            all_stroke_rgbas.extend_from_slice(&vitem.stroke_rgbas);
152            all_stroke_widths.extend_from_slice(&vitem.stroke_widths);
153
154            point_offset += pc;
155            attr_offset += ac;
156        }
157
158        // Build clip_boxes initial values: [MAX, MIN, MAX, MIN, 0] per item
159        let mut clip_boxes = Vec::with_capacity(item_count * 5);
160        for _ in 0..item_count {
161            clip_boxes.extend_from_slice(&[i32::MAX, i32::MIN, i32::MAX, i32::MIN, 0]);
162        }
163
164        // Points2d: zeroed, same size as points3d
165        let points2d = vec![Vec4::ZERO; total_points];
166
167        self.item_count = item_count as u32;
168        self.total_points = total_points as u32;
169
170        // Upload all data — track if any buffer was reallocated
171        let mut any_realloc = false;
172        any_realloc |= self.item_infos_buffer.set(ctx, &item_infos);
173        any_realloc |= self.planes_buffer.set(ctx, &planes);
174        any_realloc |= self.clip_boxes_buffer.set(ctx, &clip_boxes);
175        any_realloc |= self.points3d_buffer.set(ctx, &all_points3d);
176        any_realloc |= self.points2d_buffer.set(ctx, &points2d);
177        any_realloc |= self.fill_rgbas_buffer.set(ctx, &all_fill_rgbas);
178        any_realloc |= self.stroke_rgbas_buffer.set(ctx, &all_stroke_rgbas);
179        any_realloc |= self.stroke_widths_buffer.set(ctx, &all_stroke_widths);
180
181        // Recreate bind groups if any buffer was reallocated
182        if any_realloc || self.compute_bind_group.is_none() {
183            self.compute_bind_group = Some(Self::create_compute_bind_group(ctx, self));
184            self.render_bind_group = Some(Self::create_render_bind_group(ctx, self));
185        }
186    }
187
188    pub fn item_count(&self) -> u32 {
189        self.item_count
190    }
191
192    pub fn total_points(&self) -> u32 {
193        self.total_points
194    }
195
196    // MARK: Bind group layouts
197
198    pub fn compute_bind_group_layout(ctx: &WgpuContext) -> wgpu::BindGroupLayout {
199        ctx.device
200            .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
201                label: Some("Merged VItem Compute BGL"),
202                entries: &[
203                    // binding 0: item_infos (read-only)
204                    bgl_entry(0, wgpu::ShaderStages::COMPUTE, false),
205                    // binding 1: planes (read-only)
206                    bgl_entry(1, wgpu::ShaderStages::COMPUTE, false),
207                    // binding 2: points3d (read-only)
208                    bgl_entry(2, wgpu::ShaderStages::COMPUTE, false),
209                    // binding 3: stroke_widths (read-only)
210                    bgl_entry(3, wgpu::ShaderStages::COMPUTE, false),
211                    // binding 4: points2d (read-write)
212                    bgl_entry(4, wgpu::ShaderStages::COMPUTE, true),
213                    // binding 5: clip_boxes (read-write)
214                    bgl_entry(5, wgpu::ShaderStages::COMPUTE, true),
215                ],
216            })
217    }
218
219    pub fn render_bind_group_layout(ctx: &WgpuContext) -> wgpu::BindGroupLayout {
220        let vf = wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT;
221        let v = wgpu::ShaderStages::VERTEX;
222        ctx.device
223            .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
224                label: Some("Merged VItem Render BGL"),
225                entries: &[
226                    // binding 0: item_infos
227                    bgl_entry(0, vf, false),
228                    // binding 1: planes (normal + origin)
229                    bgl_entry(1, v, false),
230                    // binding 2: clip_boxes
231                    bgl_entry(2, v, false),
232                    // binding 3: points2d
233                    bgl_entry(3, vf, false),
234                    // binding 4: fill_rgbas
235                    bgl_entry(4, vf, false),
236                    // binding 5: stroke_rgbas
237                    bgl_entry(5, vf, false),
238                    // binding 6: stroke_widths
239                    bgl_entry(6, vf, false),
240                ],
241            })
242    }
243
244    fn create_compute_bind_group(ctx: &WgpuContext, this: &Self) -> wgpu::BindGroup {
245        ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
246            label: Some("Merged VItem Compute BG"),
247            layout: &Self::compute_bind_group_layout(ctx),
248            entries: &[
249                bg_entry(0, &this.item_infos_buffer.buffer),
250                bg_entry(1, &this.planes_buffer.buffer),
251                bg_entry(2, &this.points3d_buffer.buffer),
252                bg_entry(3, &this.stroke_widths_buffer.buffer),
253                bg_entry(4, &this.points2d_buffer.buffer),
254                bg_entry(5, &this.clip_boxes_buffer.buffer),
255            ],
256        })
257    }
258
259    fn create_render_bind_group(ctx: &WgpuContext, this: &Self) -> wgpu::BindGroup {
260        ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
261            label: Some("Merged VItem Render BG"),
262            layout: &Self::render_bind_group_layout(ctx),
263            entries: &[
264                bg_entry(0, &this.item_infos_buffer.buffer),
265                bg_entry(1, &this.planes_buffer.buffer),
266                bg_entry(2, &this.clip_boxes_buffer.buffer),
267                bg_entry(3, &this.points2d_buffer.buffer),
268                bg_entry(4, &this.fill_rgbas_buffer.buffer),
269                bg_entry(5, &this.stroke_rgbas_buffer.buffer),
270                bg_entry(6, &this.stroke_widths_buffer.buffer),
271            ],
272        })
273    }
274}
275
276fn bgl_entry(
277    binding: u32,
278    visibility: wgpu::ShaderStages,
279    read_write: bool,
280) -> wgpu::BindGroupLayoutEntry {
281    wgpu::BindGroupLayoutEntry {
282        binding,
283        visibility,
284        ty: wgpu::BindingType::Buffer {
285            ty: wgpu::BufferBindingType::Storage {
286                read_only: !read_write,
287            },
288            has_dynamic_offset: false,
289            min_binding_size: None,
290        },
291        count: None,
292    }
293}
294
295fn bg_entry(binding: u32, buffer: &wgpu::Buffer) -> wgpu::BindGroupEntry<'_> {
296    wgpu::BindGroupEntry {
297        binding,
298        resource: wgpu::BindingResource::Buffer(buffer.as_entire_buffer_binding()),
299    }
300}