Skip to main content

wde_renderer/assets/
mesh.rs

1use wde_logger::prelude::*;
2
3use bevy::{
4    asset::{AssetLoader, LoadContext, io::Reader},
5    ecs::system::{
6        SystemParamItem,
7        lifetimeless::{SRes, SResMut}
8    },
9    platform::collections::HashMap,
10    prelude::*
11};
12use serde::{Deserialize, Serialize};
13use std::{
14    fs::File,
15    io::{BufReader, Error}
16};
17use thiserror::Error;
18use tobj::LoadError;
19use wde_wgpu::{
20    buffer::{Buffer, BufferUsage},
21    vertex::Vertex
22};
23
24use crate::{
25    assets::{GpuBuffer, RenderAssets, SRenderData},
26    core::RenderInstance,
27    utils::{SsboMesh, ssbo_mesh::SsboMeshDescriptor}
28};
29
30use super::asset::{PrepareAssetError, RenderAsset};
31
32/// Utils component that stores a [`Mesh`] asset handle for 3D rendering.
33#[derive(Component, Reflect, Default, Clone)]
34#[reflect(Component)]
35pub struct Mesh3d(pub Handle<Mesh>);
36
37#[derive(Clone, Debug)]
38pub struct MeshBbox {
39    pub min: Vec3,
40    pub max: Vec3
41}
42/// Stores a CPU mesh with vertex and index data, along with metadata for GPU allocation.
43/// If loaded from a file, the mesh should have a `.obj` or `.fbx` extension.
44/// The `vertices`, `indices`, and `bbox` fields are expected to be populated by the asset loader; if not loaded from a file, they will default to empty and a degenerate bounding box respectively.
45/// This mesh will be uploaded to the GPU, represented by a [`GpuMesh`] asset.
46#[derive(Asset, TypePath, Clone, Debug)]
47pub struct Mesh {
48    pub label: String,
49
50    pub vertices: Vec<Vertex>,
51    pub indices: Vec<u32>,
52    pub bbox: MeshBbox,
53
54    /// If true, the vertices and indices will automatically be placed in the SSBO mesh buffers [`SsboMesh`] for GPU access; if false, separate vertex and index buffers will be created with `VERTEX` and `INDEX` usage respectively. Defaults to true.
55    pub use_ssbo: bool
56}
57
58/// Represents a GPU mesh resource allocated from a CPU [`Mesh`] asset.
59/// If the mesh was prepared with `use_ssbo = true`, the vertex and index data will be in the SSBO mesh buffer [`SsboMesh`] and the `vertex_buffer` and `index_buffer` fields will be `None`. Otherwhise, they will contain GPU buffers with the vertex and index data respectively.
60pub struct GpuMesh {
61    pub label: String,
62
63    /// Offset to the vertex buffer in the [`SsboMesh`] if `use_ssbo` is true. 0 otherwise.
64    pub ssbo_first_vertex: u32,
65    /// Offset to the index buffer in the [`SsboMesh`] if `use_ssbo` is true. 0 otherwise.
66    pub ssbo_first_index: u32,
67    pub use_ssbo: bool,
68
69    /// If not using SSBO, the GPU vertex buffer containing `Vertex` data. `None` if `use_ssbo` is true.
70    pub vertex_buffer: Option<Buffer>,
71    /// If not using SSBO, the GPU index buffer containing `u32` index data. `None` if `use_ssbo` is true.
72    pub index_buffer: Option<Buffer>,
73
74    pub index_count: u32,
75    pub bounding_box: MeshBbox
76}
77
78/// The SSBO region a given CPU [`Mesh`] asset currently occupies.
79struct MeshSsboAllocation {
80    first_vertex: u32,
81    first_index: u32,
82    vertex_count: u32,
83    index_count: u32
84}
85
86/// Tracks the SSBO region owned by each [`Mesh`] asset, keyed by its stable asset ID.
87/// When a mesh is re-prepared (e.g. after its vertex data is edited in place) with the same vertex and
88/// index count as before, its existing region is reused and overwritten instead of bump-allocating a new
89/// one from [`SsboMeshDescriptor`], which never reclaims freed regions.
90#[derive(Resource, Default)]
91pub struct MeshSsboAllocations(HashMap<AssetId<Mesh>, MeshSsboAllocation>);
92
93impl RenderAsset for GpuMesh {
94    type SourceAsset = Mesh;
95    type Params = (
96        SRes<RenderInstance>,
97        SResMut<SsboMeshDescriptor>,
98        SRenderData<SsboMesh>,
99        SRes<RenderAssets<GpuBuffer>>,
100        SResMut<MeshSsboAllocations>
101    );
102
103    fn prepare(
104        id: AssetId<Self::SourceAsset>,
105        asset: Self::SourceAsset,
106        (render_instance, ssbo_descriptor, ssbo_mesh, gpu_buffers, mesh_allocations): &mut SystemParamItem<
107            Self::Params
108        >
109    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
110        trace!(asset.label, "Preparing GPU mesh asset.");
111
112        // Get the SSBO mesh resource
113        let ssbo = match ssbo_mesh.iter().next() {
114            Some((_, ssbo)) => ssbo,
115            None => return Err(PrepareAssetError::RetryNextUpdate(asset))
116        };
117
118        // Get the ssbo buffers
119        let (ssbo_vertex_buffer, ssbo_index_buffer) = match (
120            gpu_buffers.get(&ssbo.get_buffer(SsboMesh::VERTEX_BUFFER_ID).unwrap()),
121            gpu_buffers.get(&ssbo.get_buffer(SsboMesh::INDEX_BUFFER_ID).unwrap())
122        ) {
123            (Some(vb), Some(ib)) => (vb, ib),
124            _ => return Err(PrepareAssetError::RetryNextUpdate(asset))
125        };
126
127        // Buffer usage
128        let usage_vertex = if asset.use_ssbo {
129            BufferUsage::COPY_SRC
130        } else {
131            BufferUsage::VERTEX
132        };
133        let usage_index = if asset.use_ssbo {
134            BufferUsage::COPY_SRC
135        } else {
136            BufferUsage::INDEX
137        };
138
139        // Create staging buffers
140        let render_instance = render_instance.0.read().unwrap();
141        let vertex_buffer = Buffer::new(
142            &render_instance,
143            format!("{}-vertex-staging", asset.label).as_str(),
144            std::mem::size_of::<Vertex>() * asset.vertices.len(),
145            usage_vertex,
146            Some(bytemuck::cast_slice(&asset.vertices))
147        );
148        let index_buffer = Buffer::new(
149            &render_instance,
150            format!("{}-indices-staging", asset.label).as_str(),
151            std::mem::size_of::<u32>() * asset.indices.len(),
152            usage_index,
153            Some(bytemuck::cast_slice(&asset.indices))
154        );
155
156        // If not using SSBO, return the buffers directly
157        if !asset.use_ssbo {
158            return Ok(GpuMesh {
159                label: asset.label,
160                ssbo_first_vertex: 0,
161                ssbo_first_index: 0,
162                index_count: asset.indices.len() as u32,
163                bounding_box: asset.bbox,
164                use_ssbo: asset.use_ssbo,
165                vertex_buffer: Some(vertex_buffer),
166                index_buffer: Some(index_buffer)
167            });
168        }
169
170        // Copy to GPU buffers, reusing this asset's existing SSBO region if its vertex/index count
171        // hasn't changed since last time, instead of bump-allocating a new one
172        let vertices_count = asset.vertices.len() as u32;
173        let indices_count = asset.indices.len() as u32;
174
175        let reused = mesh_allocations.0.get(&id).filter(|alloc| {
176            alloc.vertex_count == vertices_count && alloc.index_count == indices_count
177        });
178        let (first_vertex, first_index) = match reused {
179            Some(alloc) => (alloc.first_vertex, alloc.first_index),
180            None => {
181                let first_vertex = ssbo_descriptor.vertex_buffer_offset;
182                let first_index = ssbo_descriptor.index_buffer_offset;
183                ssbo_descriptor.vertex_buffer_offset += vertices_count;
184                ssbo_descriptor.index_buffer_offset += indices_count;
185                mesh_allocations.0.insert(
186                    id,
187                    MeshSsboAllocation {
188                        first_vertex,
189                        first_index,
190                        vertex_count: vertices_count,
191                        index_count: indices_count
192                    }
193                );
194                (first_vertex, first_index)
195            }
196        };
197
198        // Calculate byte offsets and sizes for buffer copy operations
199        let vertices_offset_bytes = (first_vertex as u64) * (std::mem::size_of::<Vertex>() as u64);
200        let indices_offset_bytes = (first_index as u64) * (std::mem::size_of::<u32>() as u64);
201        let vertices_size_bytes = (vertices_count as u64) * (std::mem::size_of::<Vertex>() as u64);
202        let indices_size_bytes = (indices_count as u64) * (std::mem::size_of::<u32>() as u64);
203
204        ssbo_vertex_buffer.buffer.copy_from_buffer_offset(
205            &render_instance,
206            &vertex_buffer,
207            0,
208            vertices_offset_bytes,
209            vertices_size_bytes
210        );
211
212        ssbo_index_buffer.buffer.copy_from_buffer_offset(
213            &render_instance,
214            &index_buffer,
215            0,
216            indices_offset_bytes,
217            indices_size_bytes
218        );
219
220        Ok(GpuMesh {
221            label: asset.label,
222            ssbo_first_vertex: first_vertex,
223            ssbo_first_index: first_index,
224            index_count: indices_count,
225            bounding_box: asset.bbox,
226            use_ssbo: asset.use_ssbo,
227            vertex_buffer: None,
228            index_buffer: None
229        })
230    }
231
232    fn label(&self) -> &str {
233        &self.label
234    }
235}
236
237/// Settings for loading a mesh asset while creating a [`Mesh`].
238#[derive(Serialize, Deserialize)]
239pub struct MeshLoaderSettings {
240    pub label: String,
241    /// Should the vertices and indices be in the SSBO mesh buffers? Defaults to true. If false, separate vertex and index buffers will be created.
242    pub use_ssbo: bool
243}
244impl Default for MeshLoaderSettings {
245    fn default() -> Self {
246        Self {
247            label: "Unknown Mesh".to_string(),
248            use_ssbo: true
249        }
250    }
251}
252
253#[derive(Debug, Error)]
254pub(crate) enum MeshLoaderError {
255    #[error("Could not load mesh: {0}")]
256    Io(#[from] std::io::Error)
257}
258#[derive(Default, TypePath)]
259pub(crate) struct MeshLoader;
260impl AssetLoader for MeshLoader {
261    type Asset = Mesh;
262    type Settings = MeshLoaderSettings;
263    type Error = MeshLoaderError;
264
265    async fn load(
266        &self,
267        reader: &mut dyn Reader,
268        settings: &Self::Settings,
269        load_context: &mut LoadContext<'_>
270    ) -> Result<Self::Asset, Self::Error> {
271        info!("Loading mesh {}.", load_context.path());
272
273        // Update the label from the path
274        let label = if settings.label.is_empty() {
275            load_context.path().to_string()
276        } else {
277            settings.label.clone()
278        };
279
280        // Read the texture data
281        let mut bytes = Vec::new();
282        reader.read_to_end(&mut bytes).await?;
283
284        // Open file
285        let load_res = match tobj::load_obj_buf(
286            &mut BufReader::new(bytes.as_slice()),
287            &tobj::LoadOptions {
288                single_index: true,
289                ..Default::default()
290            },
291            |p| {
292                let f = match File::open(p.file_name().unwrap().to_str().unwrap()) {
293                    Ok(f) => f,
294                    Err(_) => return Err(LoadError::OpenFileFailed)
295                };
296                tobj::load_mtl_buf(&mut BufReader::new(f))
297            }
298        ) {
299            Ok(res) => res,
300            Err(e) => return Err(MeshLoaderError::Io(Error::other(e.to_string())))
301        };
302        let models = load_res.0;
303
304        // Load models
305        let mut vertices: Vec<Vertex> = Vec::new();
306        let mut indices: Vec<u32> = Vec::new();
307        let mut bounding_box = MeshBbox {
308            min: Vec3::new(f32::MAX, f32::MAX, f32::MAX),
309            max: Vec3::new(f32::MIN, f32::MIN, f32::MIN)
310        };
311        for m in models.iter() {
312            let mesh = &m.mesh;
313            if mesh.positions.len() % 3 != 0 {
314                return Err(MeshLoaderError::Io(std::io::Error::other(
315                    "Mesh positions are not divisible by 3."
316                )));
317            }
318
319            // Allocate sizes
320            vertices.reserve(mesh.positions.len() / 3);
321
322            // Create vertices
323            for vtx in 0..mesh.positions.len() / 3 {
324                let x = mesh.positions[3 * vtx];
325                let y = mesh.positions[3 * vtx + 1];
326                let z = mesh.positions[3 * vtx + 2];
327
328                // Normals
329                let mut nx = 0.0;
330                let mut ny = 0.0;
331                let mut nz = 0.0;
332                if mesh.normals.len() >= 3 * vtx + 2 {
333                    nx = mesh.normals[3 * vtx];
334                    ny = mesh.normals[3 * vtx + 1];
335                    nz = mesh.normals[3 * vtx + 2];
336                }
337
338                // UVs
339                let mut u = 0.0;
340                let mut v = 0.0;
341                if mesh.texcoords.len() > 2 * vtx {
342                    u = mesh.texcoords[2 * vtx];
343                    v = mesh.texcoords[2 * vtx + 1];
344                }
345
346                // Vertex
347                vertices.push(Vertex {
348                    position: [x, y, z],
349                    normal: [nx, ny, nz],
350                    uv: [u, v],
351                    tangent: [0.0, 0.0, 0.0, 0.0]
352                });
353
354                // Update bounding box
355                bounding_box.min.x = bounding_box.min.x.min(x);
356                bounding_box.min.y = bounding_box.min.y.min(y);
357                bounding_box.min.z = bounding_box.min.z.min(z);
358                bounding_box.max.x = bounding_box.max.x.max(x);
359                bounding_box.max.y = bounding_box.max.y.max(y);
360                bounding_box.max.z = bounding_box.max.z.max(z);
361            }
362
363            // Push indices
364            indices.extend_from_slice(&mesh.indices);
365        }
366        Ok(Mesh {
367            label,
368            vertices,
369            indices,
370            bbox: bounding_box,
371            use_ssbo: settings.use_ssbo
372        })
373    }
374
375    fn extensions(&self) -> &[&str] {
376        &["obj", "fbx"]
377    }
378}