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