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#[repr(C)]
12#[derive(Debug, Default, Clone, Copy, Pod, Zeroable)]
13pub struct ItemInfo {
14 pub point_offset: u32,
16 pub point_count: u32,
18 pub attr_offset: u32,
20 pub attr_count: u32,
22}
23
24#[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#[repr(C)]
44#[derive(Debug, Default, Clone, Copy, Pod, Zeroable)]
45pub struct PlaneData {
46 pub normal: Vec4, pub origin: Vec4, }
49
50#[derive(bevy_ecs::prelude::Resource)]
56pub struct VItemsBuffer {
57 pub(crate) item_infos_buffer: WgpuVecBuffer<ItemInfo>,
59 pub(crate) planes_buffer: WgpuVecBuffer<PlaneData>,
61 pub(crate) clip_boxes_buffer: WgpuVecBuffer<i32>,
63 pub(crate) transforms_buffer: WgpuVecBuffer<VitemTransform>,
65
66 pub(crate) points3d_buffer: WgpuVecBuffer<Vec4>,
68 pub(crate) points2d_buffer: WgpuVecBuffer<Vec4>,
70 pub(crate) fill_rgbas_buffer: WgpuVecBuffer<Rgba>,
72 pub(crate) stroke_rgbas_buffer: WgpuVecBuffer<Rgba>,
74 pub(crate) stroke_widths_buffer: WgpuVecBuffer<Width>,
76
77 pub(crate) item_count: u32,
79 pub(crate) total_points: u32,
81
82 pub(crate) compute_bind_group: Option<wgpu::BindGroup>,
84 pub(crate) render_bind_group: Option<wgpu::BindGroup>,
86}
87
88impl VItemsBuffer {
89 pub fn new(ctx: &WgpuContext) -> Self {
90 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 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 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 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: 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 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 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 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 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 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 bgl_entry(0, wgpu::ShaderStages::COMPUTE, false),
239 bgl_entry(1, wgpu::ShaderStages::COMPUTE, false),
241 bgl_entry(2, wgpu::ShaderStages::COMPUTE, false),
243 bgl_entry(3, wgpu::ShaderStages::COMPUTE, false),
245 bgl_entry(4, wgpu::ShaderStages::COMPUTE, true),
247 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 bgl_entry(0, vf, false),
262 bgl_entry(1, vf, false),
264 bgl_entry(2, v, false),
266 bgl_entry(3, vf, false),
268 bgl_entry(4, vf, false),
270 bgl_entry(5, vf, false),
272 bgl_entry(6, vf, false),
274 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 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 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}