rust-engine/src/vulkan/mesh.rs

use std::{collections::HashMap, io::Read, time::SystemTime};

use gltf::{Document, Error, buffer::Source};
use gltf::mesh::util::{ReadJoints, ReadNormals, ReadPositions, ReadTangents, ReadTexCoords, ReadWeights};

use crate::vulkan::mesh::LoadError::{GltfError, MeshDataMissing, NoIndices};
use crate::vulkan::Vertex;

use super::TextVertex;

#[derive(Debug)]
pub enum LoadError {
    GltfError(gltf::Error),
    MeshDataMissing(String),
    NoIndices
}

impl From<gltf::Error> for LoadError {
    fn from(e: Error) -> Self {
        GltfError(e)
    }
}

impl From<String> for LoadError {
    fn from(e: String) -> Self {
        MeshDataMissing(e)
    }
}

#[derive(Debug)]
pub enum CPUVertexList {
    Vertex3D(Vec<Vertex>),
    VertexText(Vec<TextVertex>)
}

#[derive(Debug)]
pub struct CPUMesh {
    pub vertices: CPUVertexList,
    pub indices: Vec<u32>,
    pub local_texture_index: Option<usize>,
    pub local_normal_map_index: Option<usize>,
    pub name: Option<String>,
}

fn read_file(path: &str) -> Vec<u8> {
    let mut glb_file = std::fs::File::open(path).unwrap();
    let mut glb_bytes = vec![];
    glb_file.read_to_end(&mut glb_bytes).unwrap();
    glb_bytes
}

pub fn load_mesh<V>(mesh_path: &str, print_status: bool) -> Result<(Vec<CPUMesh>, Document), LoadError> {
    let mut start_time = None;
    let mut total_vertices = 0;
    let mut total_indices = 0;

    if print_status {
        start_time = Some(SystemTime::now());
    }
    let document = gltf::Gltf::open(mesh_path)?.document;

    let mut meshes = vec![];
    let mut mesh_data: HashMap<&str, Vec<u8>> = HashMap::new();

    for b in document.buffers() {
        if let Source::Uri(uri) = b.source() {
            mesh_data.insert(uri, read_file(&format!("models/{}", uri)));
        }
    }

    for mesh in document.meshes() {
        for primitive in mesh.primitives() {
            let texture_index = primitive.material().pbr_metallic_roughness().base_color_texture().map(|tex_info| tex_info.texture().index());
            let normal_map_index = primitive.material().normal_texture().map(|tex_info| tex_info.texture().index());

            let reader = primitive.reader(|buffer| {
                if let Source::Uri(uri) = buffer.source() {
                    Some(&mesh_data[uri])
                } else {
                    panic!("Loading embedded gltf binary is not supported!");
                }
            });
            let indices = reader.read_indices().ok_or(NoIndices)?;
            let vertices_result = create_vertices(
                reader.read_positions(),
                reader.read_tex_coords(0),
                reader.read_normals(),
                reader.read_tangents(),
                reader.read_joints(0),
                reader.read_weights(0));
            
            let verts = vertices_result?;
            let vert_count = verts.len();

            let cpu_mesh = CPUMesh {
                vertices: CPUVertexList::Vertex3D(verts),
                indices: indices.into_u32().collect(),
                local_texture_index: texture_index,
                local_normal_map_index: normal_map_index,
                name: mesh.name().map(|n| n.to_owned()),
            };
            if print_status {
                let index_count = cpu_mesh.indices.len();

                total_vertices += vert_count;
                total_indices += index_count;
            }
            meshes.push(cpu_mesh);
        }
    }

    if print_status {
        println!("Finished loading {} after {}ms, total vertices: {}, total indices: {}", mesh_path, start_time.unwrap().elapsed().unwrap().as_millis(), total_vertices, total_indices);
    }

    Ok((meshes, document))
}

fn create_vertices(positions: Option<ReadPositions>,
                   tex_coords: Option<ReadTexCoords>,
                   normals: Option<ReadNormals>,
                   tangents: Option<ReadTangents>,
                   joints: Option<ReadJoints>,
                   weights: Option<ReadWeights>)
    -> Result<Vec<Vertex>, String> {
    match (positions, tex_coords, normals, tangents, joints, weights) {
        (Some(positions),
            Some(tex_coords),
            Some(normals),
            Some(tangents),
            Some(joints),
            Some(weights)) => {
            Ok(positions
                .zip(tex_coords.into_f32())
                .zip(normals)
                .zip(tangents)
                .zip(joints.into_u16().map(|arr| {
                    let mut casted_joints: [i32; 4] = [0; 4];
                    for i in 0..4 {
                        casted_joints[i] = arr[i] as i32;
                    }
                    casted_joints
                }))
                .zip(weights.into_f32())
                .map(|(((((p, t), n), ta), i), w)| Vertex {
                    position: p,
                    uv: t,
                    normal: n,
                    tangent: ta,
                    bone_index: i,
                    bone_weight: w
                }).collect())
        },
        (Some(positions),
            Some(tex_coords),
            Some(normals),
            Some(tangents),
            None, None) => {
            Ok(positions
                .zip(tex_coords.into_f32())
                .zip(normals)
                .zip(tangents)
                .map(|(((p, t), n), ta)| Vertex {
                    position: p,
                    uv: t,
                    normal: n,
                    tangent: ta,
                    bone_index: [-1; 4],
                    bone_weight: [0.0; 4]
                }).collect())
        },
        (None, _, _, _, _, _) => Err("Vertex positions missing!".to_string()),
        (_, None, _, _, _, _) => Err("Tex coords missing!".to_string()),
        (_, _, None, _, _, _) => Err("Normals missing!".to_string()),
        (_, _, _, None, _, _) => Err("Tangents missing!".to_string()),
        (_, _, _, _, Some(_), None) => Err("Bone indices exist, but bone weights are missing!".to_string()),
        (_, _, _, _, None, Some(_)) => Err("Bone weights exist, but bone incides are missing!".to_string()),
    }
}