Skip to main content

ranim_items/vitem/render/
buffer.rs

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