zenyx-engine/src/model/mod.rs

use std::io::{BufRead, BufReader};

use crate::texture::Texture;
use cgmath::{Vector2, Vector3, Zero};
use tobj::Model as tModel;
use wgpu::util::DeviceExt;

pub struct Model {
    pub meshes: Vec<Mesh>,
    pub materials: Vec<Material>,
}

pub struct Material {
    pub name: String,
    pub diffuse_texture: Texture,
    pub bind_group: wgpu::BindGroup,
}

pub struct Mesh {
    pub name: String,
    pub vertex_buffer: wgpu::Buffer,
    pub index_buffer: wgpu::Buffer,
    pub num_elements: u32,
    pub material: Option<usize>,
}

impl Model {
    pub fn load_obj<R: BufRead>(
        read: &mut R,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        layout: &wgpu::BindGroupLayout,
    ) -> Self {
        let (models, obj_materials) = tobj::load_obj_buf(
            read,
            &tobj::LoadOptions {
                triangulate: true,
                single_index: true,
                ..Default::default()
            },
            |p| tobj::load_mtl_buf(&mut BufReader::new(std::fs::File::open(p).unwrap())),
        )
        .unwrap();

        let mut materials = Vec::new();

        for m in obj_materials.unwrap() {
            let Some(texture) = &m.diffuse_texture else {
                continue;
            };
            let diffuse_texture = load_texture(&m.name,  texture, device, queue);
            let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("model-texture-bind"),
                layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: wgpu::BindingResource::TextureView(&diffuse_texture.view),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: wgpu::BindingResource::Sampler(&diffuse_texture.sampler),
                    },
                ],
            });

            materials.push(Material {
                name: m.name,
                diffuse_texture,
                bind_group,
            });
        }

        let meshes = models
            .into_iter()
            .map(|m| {
                let vertices = (0..m.mesh.positions.len() / 3)
                    .map(|i| {
                        let position = cgmath::vec3(
                            m.mesh.positions[i * 3],
                            m.mesh.positions[i * 3 + 1],
                            m.mesh.positions[i * 3 + 2],
                        );
                        let tex_coords = cgmath::vec2(
                            m.mesh.texcoords[i * 2],
                            1.0 - m.mesh.texcoords[i * 2 + 1],
                        );
                        if m.mesh.normals.is_empty() {
                            Vertex {
                                position,
                                tex_coords,
                                normal: cgmath::Vector3::zero(),
                                color: cgmath::Vector3::zero(),
                            }
                        } else {
                            Vertex {
                                position,
                                tex_coords,
                                normal: cgmath::vec3(
                                    m.mesh.normals[i * 3],
                                    m.mesh.normals[i * 3 + 1],
                                    m.mesh.normals[i * 3 + 2],
                                ),
                                color: cgmath::Vector3::zero(),
                            }
                        }
                    })
                    .collect::<Vec<_>>();

                let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytemuck::cast_slice(&vertices),
                    usage: wgpu::BufferUsages::VERTEX,
                });
                let index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytemuck::cast_slice(&m.mesh.indices),
                    usage: wgpu::BufferUsages::INDEX,
                });

                Mesh {
                    name: "TOBJ_MESH".to_string(),
                    vertex_buffer,
                    index_buffer,
                    num_elements: m.mesh.indices.len() as u32,
                    material: m.mesh.material_id,
                }
            })
            .collect::<Vec<_>>();

        Self { meshes, materials }
    }
}

pub fn load_texture<S: AsRef<std::path::Path>>(
    label: &str,
    path: S,
    device: &wgpu::Device,
    queue: &wgpu::Queue,
) -> Texture {
    let file = std::fs::read(path).unwrap();
    Texture::from_bytes(device, queue, &file, label).unwrap()
}

pub fn load(obj: &[tModel]) -> (Vec<Vertex>, Vec<u32>) {
    let mut combined_vertices = Vec::new();
    let mut combined_indices = Vec::new();
    let mut vertex_offset = 0;

    for object in obj {
        let mesh: &_ = &object.mesh;
        let vertices: Vec<Vertex> = (0..mesh.positions.len() / 3)
            .map(|i| Vertex {
                position: Vector3::from([
                    mesh.positions[i * 3],
                    mesh.positions[i * 3 + 1],
                    mesh.positions[i * 3 + 2],
                ]),
                color: cgmath::Vector3::from([1.0, 1.0, 1.0]),
                normal: if !mesh.normals.is_empty() {
                    Vector3::from([
                        mesh.normals[i * 3],
                        mesh.normals[i * 3 + 1],
                        mesh.normals[i * 3 + 2],
                    ])
                } else {
                    Vector3::zero()
                },
                tex_coords: if !mesh.texcoords.is_empty() {
                    Vector2::from([mesh.texcoords[i * 2], mesh.texcoords[i * 2 + 1]])
                } else {
                    Vector2::zero()
                },
            })
            .collect();
        combined_vertices.extend(vertices);
        combined_indices.extend(mesh.indices.iter().map(|&index| index + vertex_offset));
        vertex_offset += (mesh.positions.len() as u32) / 3;
    }

    (combined_vertices, combined_indices)
}

#[derive(Copy, Clone, Debug)]
#[repr(C)]
pub struct Vertex {
    pub position: cgmath::Vector3<f32>,
    pub color: cgmath::Vector3<f32>,
    pub normal: cgmath::Vector3<f32>,
    pub tex_coords: cgmath::Vector2<f32>,
}

unsafe impl bytemuck::Pod for Vertex {}
unsafe impl bytemuck::Zeroable for Vertex {}
impl Vertex {
    pub const ATTRIBS: [wgpu::VertexAttribute; 4] = [
        wgpu::VertexAttribute {
            offset: 0,
            shader_location: 0,
            format: wgpu::VertexFormat::Float32x3,
        },
        wgpu::VertexAttribute {
            offset: std::mem::offset_of!(Vertex, color) as u64,
            shader_location: 1,
            format: wgpu::VertexFormat::Float32x3,
        },
        wgpu::VertexAttribute {
            offset: std::mem::offset_of!(Vertex, normal) as u64,
            shader_location: 2,
            format: wgpu::VertexFormat::Float32x3,
        },
        wgpu::VertexAttribute {
            offset: std::mem::offset_of!(Vertex, tex_coords) as u64,
            shader_location: 3,
            format: wgpu::VertexFormat::Float32x2,
        },
    ];

    pub const fn desc<'a>() -> wgpu::VertexBufferLayout<'a> {
        wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &Self::ATTRIBS,
        }
    }
}