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 local-to-world transform, applied by the vertex stage after
25/// reconstructing the 3D position from the plane basis.
26#[repr(C)]
27#[derive(Debug, Clone, Copy, Pod, Zeroable)]
28pub struct VitemTransform {
29    pub transform: [[f32; 4]; 4],
30}
31
32impl Default for VitemTransform {
33    fn default() -> Self {
34        Self {
35            transform: glam::Mat4::IDENTITY.to_cols_array_2d(),
36        }
37    }
38}
39
40/// Per-item plane data (normal + origin), stored as array of structs.
41/// The origin is the first point of the item (used by vertex shader).
42/// basis_u/basis_v are generated deterministically from the normal in the shader.
43#[repr(C)]
44#[derive(Debug, Default, Clone, Copy, Pod, Zeroable)]
45pub struct PlaneData {
46    pub normal: Vec4, // xyz = normal, w = pad
47    pub origin: Vec4, // xyz = first point, w = pad
48}
49
50/// Merged GPU buffers for all VItems in a frame.
51///
52/// Instead of one set of buffers per VItem, all data is packed into
53/// contiguous arrays with an index table (`item_infos`) that tells
54/// shaders where each item's data lives.
55#[derive(bevy_ecs::prelude::Resource)]
56pub struct VItemsBuffer {
57    /// Per-item metadata: offsets and counts
58    pub(crate) item_infos_buffer: WgpuVecBuffer<ItemInfo>,
59    /// Per-item plane data (normal + origin for vertex shader)
60    pub(crate) planes_buffer: WgpuVecBuffer<PlaneData>,
61    /// Per-item clip boxes (5 i32 each: min_x, max_x, min_y, max_y, max_w)
62    pub(crate) clip_boxes_buffer: WgpuVecBuffer<i32>,
63    /// Per-item local-to-world transforms
64    pub(crate) transforms_buffer: WgpuVecBuffer<VitemTransform>,
65
66    /// Merged 3D points from all VItems
67    pub(crate) points3d_buffer: WgpuVecBuffer<Vec4>,
68    /// Merged 2D projected points (written by compute shader)
69    pub(crate) points2d_buffer: WgpuVecBuffer<Vec4>,
70    /// Merged fill colors
71    pub(crate) fill_rgbas_buffer: WgpuVecBuffer<Rgba>,
72    /// Merged stroke colors
73    pub(crate) stroke_rgbas_buffer: WgpuVecBuffer<Rgba>,
74    /// Merged stroke widths
75    pub(crate) stroke_widths_buffer: WgpuVecBuffer<Width>,
76
77    /// Number of items
78    pub(crate) item_count: u32,
79    /// Total number of points across all items
80    pub(crate) total_points: u32,
81
82    /// Compute bind group (recreated when buffers resize)
83    pub(crate) compute_bind_group: Option<wgpu::BindGroup>,
84    /// Render bind group (recreated when buffers resize)
85    pub(crate) render_bind_group: Option<wgpu::BindGroup>,
86}
87
88impl VItemsBuffer {
89    pub fn new(ctx: &WgpuContext) -> Self {
90        // Start with empty buffers (minimum size 1 to avoid zero-size buffer)
91        let storage_rw = wgpu::BufferUsages::STORAGE
92            | wgpu::BufferUsages::COPY_DST
93            | wgpu::BufferUsages::COPY_SRC;
94        let storage_ro = wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST;
95
96        Self {
97            item_infos_buffer: WgpuVecBuffer::new(ctx, Some("Merged ItemInfos"), storage_ro, 1),
98            planes_buffer: WgpuVecBuffer::new(ctx, Some("Merged Planes"), storage_ro, 1),
99            clip_boxes_buffer: WgpuVecBuffer::new(ctx, Some("Merged ClipBoxes"), storage_rw, 5),
100            transforms_buffer: WgpuVecBuffer::new(
101                ctx,
102                Some("Merged VItem Transforms"),
103                storage_ro,
104                1,
105            ),
106            points3d_buffer: WgpuVecBuffer::new(ctx, Some("Merged Points3D"), storage_ro, 1),
107            points2d_buffer: WgpuVecBuffer::new(ctx, Some("Merged Points2D"), storage_rw, 1),
108            fill_rgbas_buffer: WgpuVecBuffer::new(ctx, Some("Merged FillRgbas"), storage_ro, 1),
109            stroke_rgbas_buffer: WgpuVecBuffer::new(ctx, Some("Merged StrokeRgbas"), storage_ro, 1),
110            stroke_widths_buffer: WgpuVecBuffer::new(
111                ctx,
112                Some("Merged StrokeWidths"),
113                storage_ro,
114                1,
115            ),
116            item_count: 0,
117            total_points: 0,
118            compute_bind_group: None,
119            render_bind_group: None,
120        }
121    }
122
123    /// Pack all VItems into the merged buffers. Called once per frame.
124    pub fn update<'a, I>(&mut self, ctx: &WgpuContext, vitems: I)
125    where
126        I: IntoIterator<Item = (f32, &'a VItem)>,
127        I::IntoIter: ExactSizeIterator + Clone,
128    {
129        let vitems = vitems.into_iter();
130        if vitems.len() == 0 {
131            self.item_count = 0;
132            self.total_points = 0;
133            return;
134        }
135
136        let item_count = vitems.len();
137
138        // Pre-calculate total sizes
139        let total_points: usize = vitems.clone().map(|(_, v)| v.points.len()).sum();
140        let total_attrs: usize = vitems
141            .clone()
142            .map(|(_, v)| v.points.len().div_ceil(2))
143            .sum();
144
145        // Build index table and collect data
146        let mut item_infos = Vec::with_capacity(item_count);
147        let mut planes = Vec::with_capacity(item_count);
148        let mut transforms = Vec::with_capacity(item_count);
149        let mut all_points3d = Vec::with_capacity(total_points);
150        let mut all_fill_rgbas = Vec::with_capacity(total_attrs);
151        let mut all_stroke_rgbas = Vec::with_capacity(total_attrs);
152        let mut all_stroke_widths = Vec::with_capacity(total_attrs);
153
154        let mut point_offset: u32 = 0;
155        let mut attr_offset: u32 = 0;
156
157        for (order, vitem) in vitems {
158            let pc = vitem.points.len() as u32;
159            let ac = pc.div_ceil(2);
160
161            item_infos.push(ItemInfo {
162                point_offset,
163                point_count: pc,
164                attr_offset,
165                attr_count: ac,
166            });
167
168            let normal = vitem
169                .normal
170                .unwrap_or_else(|| vitem_normal_from_points(&vitem.points));
171            let origin = Vec3::new(vitem.points[0].x, vitem.points[0].y, vitem.points[0].z);
172            planes.push(PlaneData {
173                normal: Vec4::from((normal, 0.0)),
174                // origin.w carries the item's global scene order for the
175                // depth-order bias; the plane basis ignores it.
176                origin: Vec4::from((origin, order)),
177            });
178
179            all_points3d.extend_from_slice(&vitem.points);
180            transforms.push(VitemTransform {
181                transform: vitem.transform.to_cols_array_2d(),
182            });
183            all_fill_rgbas.extend_from_slice(&vitem.fill_rgbas);
184            all_stroke_rgbas.extend_from_slice(&vitem.stroke_rgbas);
185            all_stroke_widths.extend_from_slice(&vitem.stroke_widths);
186
187            point_offset += pc;
188            attr_offset += ac;
189        }
190
191        // Build clip_boxes initial values: [MAX, MIN, MAX, MIN, 0] per item
192        let mut clip_boxes = Vec::with_capacity(item_count * 5);
193        for _ in 0..item_count {
194            clip_boxes.extend_from_slice(&[i32::MAX, i32::MIN, i32::MAX, i32::MIN, 0]);
195        }
196
197        // Points2d: zeroed, same size as points3d
198        let points2d = vec![Vec4::ZERO; total_points];
199
200        self.item_count = item_count as u32;
201        self.total_points = total_points as u32;
202
203        // Upload all data — track if any buffer was reallocated
204        let mut any_realloc = false;
205        any_realloc |= self.item_infos_buffer.set(ctx, &item_infos);
206        any_realloc |= self.planes_buffer.set(ctx, &planes);
207        any_realloc |= self.transforms_buffer.set(ctx, &transforms);
208        any_realloc |= self.clip_boxes_buffer.set(ctx, &clip_boxes);
209        any_realloc |= self.points3d_buffer.set(ctx, &all_points3d);
210        any_realloc |= self.points2d_buffer.set(ctx, &points2d);
211        any_realloc |= self.fill_rgbas_buffer.set(ctx, &all_fill_rgbas);
212        any_realloc |= self.stroke_rgbas_buffer.set(ctx, &all_stroke_rgbas);
213        any_realloc |= self.stroke_widths_buffer.set(ctx, &all_stroke_widths);
214
215        // Recreate bind groups if any buffer was reallocated
216        if any_realloc || self.compute_bind_group.is_none() {
217            self.compute_bind_group = Some(Self::create_compute_bind_group(ctx, self));
218            self.render_bind_group = Some(Self::create_render_bind_group(ctx, self));
219        }
220    }
221
222    pub fn item_count(&self) -> u32 {
223        self.item_count
224    }
225
226    pub fn total_points(&self) -> u32 {
227        self.total_points
228    }
229
230    // MARK: Bind group layouts
231
232    pub fn compute_bind_group_layout(ctx: &WgpuContext) -> wgpu::BindGroupLayout {
233        ctx.device
234            .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
235                label: Some("Merged VItem Compute BGL"),
236                entries: &[
237                    // binding 0: item_infos (read-only)
238                    bgl_entry(0, wgpu::ShaderStages::COMPUTE, false),
239                    // binding 1: planes (read-only)
240                    bgl_entry(1, wgpu::ShaderStages::COMPUTE, false),
241                    // binding 2: points3d (read-only)
242                    bgl_entry(2, wgpu::ShaderStages::COMPUTE, false),
243                    // binding 3: stroke_widths (read-only)
244                    bgl_entry(3, wgpu::ShaderStages::COMPUTE, false),
245                    // binding 4: points2d (read-write)
246                    bgl_entry(4, wgpu::ShaderStages::COMPUTE, true),
247                    // binding 5: clip_boxes (read-write)
248                    bgl_entry(5, wgpu::ShaderStages::COMPUTE, true),
249                ],
250            })
251    }
252
253    pub fn render_bind_group_layout(ctx: &WgpuContext) -> wgpu::BindGroupLayout {
254        let vf = wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT;
255        let v = wgpu::ShaderStages::VERTEX;
256        ctx.device
257            .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
258                label: Some("Merged VItem Render BGL"),
259                entries: &[
260                    // binding 0: item_infos
261                    bgl_entry(0, vf, false),
262                    // binding 1: planes (normal + origin; origin.w = scene order)
263                    bgl_entry(1, vf, false),
264                    // binding 2: clip_boxes
265                    bgl_entry(2, v, false),
266                    // binding 3: points2d
267                    bgl_entry(3, vf, false),
268                    // binding 4: fill_rgbas
269                    bgl_entry(4, vf, false),
270                    // binding 5: stroke_rgbas
271                    bgl_entry(5, vf, false),
272                    // binding 6: stroke_widths
273                    bgl_entry(6, vf, false),
274                    // binding 7: per-item local-to-world transforms
275                    bgl_entry(7, v, false),
276                ],
277            })
278    }
279
280    fn create_compute_bind_group(ctx: &WgpuContext, this: &Self) -> wgpu::BindGroup {
281        ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
282            label: Some("Merged VItem Compute BG"),
283            layout: &Self::compute_bind_group_layout(ctx),
284            entries: &[
285                bg_entry(0, &this.item_infos_buffer.buffer),
286                bg_entry(1, &this.planes_buffer.buffer),
287                bg_entry(2, &this.points3d_buffer.buffer),
288                bg_entry(3, &this.stroke_widths_buffer.buffer),
289                bg_entry(4, &this.points2d_buffer.buffer),
290                bg_entry(5, &this.clip_boxes_buffer.buffer),
291            ],
292        })
293    }
294
295    fn create_render_bind_group(ctx: &WgpuContext, this: &Self) -> wgpu::BindGroup {
296        ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
297            label: Some("Merged VItem Render BG"),
298            layout: &Self::render_bind_group_layout(ctx),
299            entries: &[
300                bg_entry(0, &this.item_infos_buffer.buffer),
301                bg_entry(1, &this.planes_buffer.buffer),
302                bg_entry(2, &this.clip_boxes_buffer.buffer),
303                bg_entry(3, &this.points2d_buffer.buffer),
304                bg_entry(4, &this.fill_rgbas_buffer.buffer),
305                bg_entry(5, &this.stroke_rgbas_buffer.buffer),
306                bg_entry(6, &this.stroke_widths_buffer.buffer),
307                bg_entry(7, &this.transforms_buffer.buffer),
308            ],
309        })
310    }
311}
312
313fn bgl_entry(
314    binding: u32,
315    visibility: wgpu::ShaderStages,
316    read_write: bool,
317) -> wgpu::BindGroupLayoutEntry {
318    wgpu::BindGroupLayoutEntry {
319        binding,
320        visibility,
321        ty: wgpu::BindingType::Buffer {
322            ty: wgpu::BufferBindingType::Storage {
323                read_only: !read_write,
324            },
325            has_dynamic_offset: false,
326            min_binding_size: None,
327        },
328        count: None,
329    }
330}
331
332fn bg_entry(binding: u32, buffer: &wgpu::Buffer) -> wgpu::BindGroupEntry<'_> {
333    wgpu::BindGroupEntry {
334        binding,
335        resource: wgpu::BindingResource::Buffer(buffer.as_entire_buffer_binding()),
336    }
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342    use crate::primitives::test_output_path;
343    use crate::{Renderer, world::RenderFrame};
344    use glam::{Mat4, Vec3};
345    use pollster::block_on;
346    use ranim_core::{
347        components::rgba::Rgba,
348        core_item::{CoreItem, camera_frame::CameraFrame},
349    };
350
351    /// A closed unit square centered at the origin in local space.
352    fn square_vitem(color: Rgba, stroke: Rgba) -> VItem {
353        let p = |x: f32, y: f32| Vec4::new(x, y, 0.0, 1.0);
354        let mid = |a: Vec4, b: Vec4| Vec4::new((a.x + b.x) * 0.5, (a.y + b.y) * 0.5, 0.0, 1.0);
355        let a0 = p(-0.5, -0.5);
356        let a1 = p(0.5, -0.5);
357        let a2 = p(0.5, 0.5);
358        let a3 = p(-0.5, 0.5);
359        VItem {
360            normal: None,
361            points: vec![
362                a0,
363                mid(a0, a1),
364                a1,
365                mid(a1, a2),
366                a2,
367                mid(a2, a3),
368                a3,
369                mid(a3, a0),
370            ],
371            transform: Mat4::IDENTITY,
372            fill_rgbas: vec![color; 4],
373            stroke_rgbas: vec![stroke; 4],
374            stroke_widths: vec![ranim_core::components::width::Width(0.02); 4],
375        }
376    }
377
378    #[test]
379    #[ignore = "requires a GPU"]
380    fn render_transformed_vitems() {
381        let ctx = block_on(WgpuContext::new());
382        let width = 800u32;
383        let height = 600u32;
384
385        let mut renderer = Renderer::new(&ctx, width, height, 8);
386        let mut render_textures = renderer.new_render_textures(&ctx);
387        let mut store = RenderFrame::new();
388
389        // The red square is rotated around Z, the blue one scaled; both keep
390        // identity local points and are placed only through `transform`.
391        let mut red = square_vitem(
392            Rgba(glam::Vec4::new(1.0, 0.0, 0.0, 0.6)),
393            Rgba(glam::Vec4::new(1.0, 1.0, 1.0, 1.0)),
394        );
395        red.transform = Mat4::from_rotation_z(std::f32::consts::FRAC_PI_4)
396            * Mat4::from_translation(Vec3::new(2.0, 0.0, 0.0));
397        let mut blue = square_vitem(
398            Rgba(glam::Vec4::new(0.0, 0.0, 1.0, 0.6)),
399            Rgba(glam::Vec4::new(1.0, 1.0, 1.0, 1.0)),
400        );
401        blue.transform =
402            Mat4::from_scale(Vec3::splat(2.0)) * Mat4::from_translation(Vec3::new(-2.5, 0.0, 0.0));
403
404        store.update(
405            [
406                ((0, 0), CoreItem::CameraFrame(CameraFrame::default())),
407                ((1, 0), CoreItem::VItem(red)),
408                ((1, 1), CoreItem::VItem(blue)),
409            ]
410            .into_iter(),
411        );
412
413        renderer.render_frame(&mut render_textures, wgpu::Color::BLACK, &store);
414        ctx.device
415            .poll(wgpu::PollType::wait_indefinitely())
416            .unwrap();
417
418        let buffer = render_textures.get_rendered_texture_img_buffer(&ctx);
419        let output_path = test_output_path("vitems_transformed_render.png");
420        buffer.save(&output_path).expect("Failed to save image");
421        assert!(output_path.exists(), "Image file should be created");
422    }
423}