Skip to main content

ranim_items/mesh/render/
buffer.rs

1use super::RenderMeshItem;
2use bytemuck::{Pod, Zeroable};
3use ranim_core::components::rgba::Rgba;
4use ranim_core::glam::Vec3;
5use ranim_render::utils::{WgpuContext, WgpuVecBuffer};
6
7#[repr(C)]
8#[derive(Debug, Default, Clone, Copy, Pod, Zeroable)]
9pub struct MeshTransform {
10    pub transform: [[f32; 4]; 4],
11}
12
13#[derive(bevy_ecs::prelude::Resource)]
14pub struct MeshItemsBuffer {
15    /// Per-vertex positions (vertex buffer)
16    pub(crate) vertices_buffer: WgpuVecBuffer<Vec3>,
17    /// Per-vertex mesh id (vertex buffer)
18    pub(crate) mesh_ids_buffer: WgpuVecBuffer<u32>,
19    /// Per-vertex colors (vertex buffer)
20    pub(crate) vertex_colors_buffer: WgpuVecBuffer<Rgba>,
21    /// Per-vertex normals (vertex buffer) — all-zero → flat shading fallback
22    pub(crate) vertex_normals_buffer: WgpuVecBuffer<Vec3>,
23    /// Merged triangle indices (index buffer)
24    pub(crate) indices_buffer: WgpuVecBuffer<u32>,
25
26    /// Per-mesh transform matrices (storage buffer, indexed by mesh_id)
27    pub(crate) transforms_buffer: WgpuVecBuffer<MeshTransform>,
28    /// Per-mesh global scene order (storage buffer, indexed by mesh_id),
29    /// consumed by the fragment depth-order bias.
30    pub(crate) orders_buffer: WgpuVecBuffer<f32>,
31
32    pub(crate) item_count: u32,
33    pub(crate) total_vertices: u32,
34    pub(crate) total_indices: u32,
35
36    pub(crate) render_bind_group: Option<wgpu::BindGroup>,
37}
38
39impl MeshItemsBuffer {
40    pub fn new(ctx: &WgpuContext) -> Self {
41        let vertex_usage = wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST;
42        let index_usage = wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST;
43        let storage_ro = wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST;
44
45        Self {
46            vertices_buffer: WgpuVecBuffer::new(ctx, Some("MeshVertices"), vertex_usage, 1),
47            mesh_ids_buffer: WgpuVecBuffer::new(ctx, Some("MeshIds"), vertex_usage, 1),
48            vertex_colors_buffer: WgpuVecBuffer::new(
49                ctx,
50                Some("MeshVertexColors"),
51                vertex_usage,
52                1,
53            ),
54            vertex_normals_buffer: WgpuVecBuffer::new(
55                ctx,
56                Some("MeshVertexNormals"),
57                vertex_usage,
58                1,
59            ),
60            indices_buffer: WgpuVecBuffer::new(ctx, Some("MeshIndices"), index_usage, 1),
61            transforms_buffer: WgpuVecBuffer::new(ctx, Some("MeshTransforms"), storage_ro, 1),
62            orders_buffer: WgpuVecBuffer::new(ctx, Some("MeshOrders"), storage_ro, 1),
63            item_count: 0,
64            total_vertices: 0,
65            total_indices: 0,
66            render_bind_group: None,
67        }
68    }
69
70    pub fn update<'a, I>(&mut self, ctx: &WgpuContext, mesh_items: I)
71    where
72        I: IntoIterator<Item = (f32, &'a RenderMeshItem)>,
73        I::IntoIter: ExactSizeIterator + Clone,
74    {
75        let mesh_items = mesh_items.into_iter();
76        if mesh_items.len() == 0 {
77            self.item_count = 0;
78            self.total_vertices = 0;
79            self.total_indices = 0;
80            return;
81        }
82
83        let item_count = mesh_items.len();
84        let total_vertices: usize = mesh_items.clone().map(|(_, m)| m.points.len()).sum();
85        let total_indices: usize = mesh_items
86            .clone()
87            .map(|(_, m)| m.triangle_indices.len())
88            .sum();
89
90        let mut transforms = Vec::with_capacity(item_count);
91        let mut all_orders = Vec::with_capacity(item_count);
92        let mut all_vertices = Vec::with_capacity(total_vertices);
93        let mut all_mesh_ids = Vec::with_capacity(total_vertices);
94        let mut all_vertex_colors = Vec::with_capacity(total_vertices);
95        let mut all_vertex_normals = Vec::with_capacity(total_vertices);
96        let mut all_indices = Vec::with_capacity(total_indices);
97
98        let mut vertex_offset: u32 = 0;
99
100        for (mesh_idx, (order, mesh)) in mesh_items.enumerate() {
101            let vc = mesh.points.len() as u32;
102
103            transforms.push(MeshTransform {
104                transform: mesh.transform.to_cols_array_2d(),
105            });
106            all_orders.push(order);
107
108            all_vertices.extend_from_slice(&mesh.points);
109            all_mesh_ids.extend(std::iter::repeat_n(mesh_idx as u32, vc as usize));
110            all_vertex_colors.extend_from_slice(&mesh.vertex_colors);
111
112            // Pad normals with zero if shorter than points (flat shading fallback)
113            let normals = &mesh.vertex_normals;
114            let normals_len = normals.len();
115            if normals_len >= vc as usize {
116                all_vertex_normals.extend_from_slice(&normals[..vc as usize]);
117            } else {
118                all_vertex_normals.extend_from_slice(normals);
119                all_vertex_normals
120                    .extend(std::iter::repeat_n(Vec3::ZERO, vc as usize - normals_len));
121            }
122
123            all_indices.extend(mesh.triangle_indices.iter().map(|&i| i + vertex_offset));
124
125            vertex_offset += vc;
126        }
127
128        self.item_count = item_count as u32;
129        self.total_vertices = total_vertices as u32;
130        self.total_indices = total_indices as u32;
131
132        // Vertex/index buffers (no bind group dependency)
133        self.vertices_buffer.set(ctx, &all_vertices);
134        self.mesh_ids_buffer.set(ctx, &all_mesh_ids);
135        self.vertex_colors_buffer.set(ctx, &all_vertex_colors);
136        self.vertex_normals_buffer.set(ctx, &all_vertex_normals);
137        self.indices_buffer.set(ctx, &all_indices);
138
139        // Storage buffers (bind group recreated on realloc)
140        let mut any_realloc = self.transforms_buffer.set(ctx, &transforms);
141        any_realloc |= self.orders_buffer.set(ctx, &all_orders);
142
143        if any_realloc || self.render_bind_group.is_none() {
144            self.render_bind_group = Some(Self::create_render_bind_group(ctx, self));
145        }
146    }
147
148    pub fn item_count(&self) -> u32 {
149        self.item_count
150    }
151
152    pub fn total_indices(&self) -> u32 {
153        self.total_indices
154    }
155
156    pub fn vertex_buffer_layouts() -> [wgpu::VertexBufferLayout<'static>; 4] {
157        [
158            // Slot 0: positions (vec3<f32>)
159            wgpu::VertexBufferLayout {
160                array_stride: std::mem::size_of::<Vec3>() as u64,
161                step_mode: wgpu::VertexStepMode::Vertex,
162                attributes: &[wgpu::VertexAttribute {
163                    format: wgpu::VertexFormat::Float32x3,
164                    offset: 0,
165                    shader_location: 0,
166                }],
167            },
168            // Slot 1: mesh_id (u32)
169            wgpu::VertexBufferLayout {
170                array_stride: std::mem::size_of::<u32>() as u64,
171                step_mode: wgpu::VertexStepMode::Vertex,
172                attributes: &[wgpu::VertexAttribute {
173                    format: wgpu::VertexFormat::Uint32,
174                    offset: 0,
175                    shader_location: 1,
176                }],
177            },
178            // Slot 2: vertex_color (vec4<f32>)
179            wgpu::VertexBufferLayout {
180                array_stride: std::mem::size_of::<Rgba>() as u64,
181                step_mode: wgpu::VertexStepMode::Vertex,
182                attributes: &[wgpu::VertexAttribute {
183                    format: wgpu::VertexFormat::Float32x4,
184                    offset: 0,
185                    shader_location: 2,
186                }],
187            },
188            // Slot 3: vertex_normal (vec3<f32>)
189            wgpu::VertexBufferLayout {
190                array_stride: std::mem::size_of::<Vec3>() as u64,
191                step_mode: wgpu::VertexStepMode::Vertex,
192                attributes: &[wgpu::VertexAttribute {
193                    format: wgpu::VertexFormat::Float32x3,
194                    offset: 0,
195                    shader_location: 3,
196                }],
197            },
198        ]
199    }
200
201    pub fn render_bind_group_layout(ctx: &WgpuContext) -> wgpu::BindGroupLayout {
202        ctx.device
203            .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
204                label: Some("MeshItems Render BGL"),
205                entries: &[
206                    // binding 0: transforms (per-mesh, vertex stage)
207                    bgl_storage_entry(0, wgpu::ShaderStages::VERTEX),
208                    // binding 1: orders (per-mesh, fragment stage, depth bias)
209                    bgl_storage_entry(1, wgpu::ShaderStages::FRAGMENT),
210                ],
211            })
212    }
213
214    fn create_render_bind_group(ctx: &WgpuContext, this: &Self) -> wgpu::BindGroup {
215        ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
216            label: Some("MeshItems Render BG"),
217            layout: &Self::render_bind_group_layout(ctx),
218            entries: &[
219                bg_entry(0, this.transforms_buffer.buffer()),
220                bg_entry(1, this.orders_buffer.buffer()),
221            ],
222        })
223    }
224}
225
226fn bgl_storage_entry(binding: u32, visibility: wgpu::ShaderStages) -> wgpu::BindGroupLayoutEntry {
227    wgpu::BindGroupLayoutEntry {
228        binding,
229        visibility,
230        ty: wgpu::BindingType::Buffer {
231            ty: wgpu::BufferBindingType::Storage { read_only: true },
232            has_dynamic_offset: false,
233            min_binding_size: None,
234        },
235        count: None,
236    }
237}
238
239fn bg_entry(binding: u32, buffer: &wgpu::Buffer) -> wgpu::BindGroupEntry<'_> {
240    wgpu::BindGroupEntry {
241        binding,
242        resource: wgpu::BindingResource::Buffer(buffer.as_entire_buffer_binding()),
243    }
244}