Oxide/engine/src/asset/gltf.rs

//! glTF 2.0 static-mesh importer.
//!
//! Loads the mesh primitives of a glTF document into engine [`Mesh`]es, reading
//! their PBR-lite [`Material`] factors and the world [`Transform`] of each
//! placement (the node hierarchy is flattened into world space). Missing
//! normals are generated; missing UVs default to zero. Animation, skinning, and
//! textures are out of scope for Stage 4.

use std::path::Path;

use crate::math::{Color, Transform, Vec2, Vec3};
use crate::render::{Material, Mesh, Vertex};

/// Errors produced while importing a glTF document.
#[derive(Debug, thiserror::Error)]
pub enum GltfError {
    /// The file could not be read or parsed as glTF.
    #[error("failed to load glTF: {0}")]
    Load(#[from] gltf::Error),

    /// A mesh primitive was missing the required `POSITION` attribute.
    #[error("glTF primitive has no POSITION attribute")]
    MissingPositions,
}

/// One imported mesh placement: geometry, material, and world transform.
pub struct GltfMesh {
    /// Optional node/mesh name from the document.
    pub name: Option<String>,
    /// The primitive's geometry.
    pub mesh: Mesh,
    /// The primitive's PBR-lite material.
    pub material: Material,
    /// World-space placement (node hierarchy flattened).
    pub transform: Transform,
}

/// An imported glTF model: a flat list of mesh placements in world space.
pub struct GltfModel {
    /// Every mesh primitive in the default scene, already placed in world space.
    pub meshes: Vec<GltfMesh>,
}

impl GltfModel {
    /// Total triangle count across all imported primitives.
    pub fn triangle_count(&self) -> usize {
        self.meshes.iter().map(|m| m.mesh.triangle_count()).sum()
    }
}

/// Imports a glTF/GLB file from `path` (external buffers are resolved relative
/// to the file).
pub fn load_gltf(path: impl AsRef<Path>) -> Result<GltfModel, GltfError> {
    let (document, buffers, _images) = gltf::import(path)?;
    build_model(&document, &buffers)
}

/// The [`AssetServer`](super::AssetServer) loader for glTF/GLB files.
///
/// Registered by default (handles `.gltf` and `.glb`), so
/// `assets.load::<GltfModel>("model.gltf")` works out of the box; it simply
/// wraps [`load_gltf`] and adapts its error into [`AssetError`].
pub struct GltfLoader;

impl super::AssetLoader for GltfLoader {
    type Asset = GltfModel;

    fn extensions(&self) -> &'static [&'static str] {
        &["gltf", "glb"]
    }

    fn load(&self, path: &Path) -> Result<GltfModel, super::AssetError> {
        load_gltf(path).map_err(|err| super::AssetError::Load {
            path: path.to_path_buf(),
            message: err.to_string(),
        })
    }
}

/// Imports a glTF/GLB document from an in-memory byte slice (buffers must be
/// embedded; used for tests and bundled assets).
pub fn load_gltf_slice(bytes: &[u8]) -> Result<GltfModel, GltfError> {
    let (document, buffers, _images) = gltf::import_slice(bytes)?;
    build_model(&document, &buffers)
}

/// Walks the default scene's node hierarchy, accumulating world transforms and
/// emitting one [`GltfMesh`] per primitive.
fn build_model(
    document: &gltf::Document,
    buffers: &[gltf::buffer::Data],
) -> Result<GltfModel, GltfError> {
    let mut meshes = Vec::new();
    let scene = document
        .default_scene()
        .or_else(|| document.scenes().next());
    if let Some(scene) = scene {
        for node in scene.nodes() {
            visit_node(&node, Transform::IDENTITY, buffers, &mut meshes)?;
        }
    }
    Ok(GltfModel { meshes })
}

fn visit_node(
    node: &gltf::Node,
    parent: Transform,
    buffers: &[gltf::buffer::Data],
    out: &mut Vec<GltfMesh>,
) -> Result<(), GltfError> {
    let world = parent.mul_transform(&node_transform(node));

    if let Some(mesh) = node.mesh() {
        for primitive in mesh.primitives() {
            let geometry = read_primitive(&primitive, buffers)?;
            out.push(GltfMesh {
                name: node.name().or_else(|| mesh.name()).map(str::to_owned),
                mesh: geometry,
                material: read_material(&primitive),
                transform: world,
            });
        }
    }

    for child in node.children() {
        visit_node(&child, world, buffers, out)?;
    }
    Ok(())
}

/// Converts a node's local TRS into an engine [`Transform`].
fn node_transform(node: &gltf::Node) -> Transform {
    let (t, r, s) = node.transform().decomposed();
    Transform::from_trs(
        Vec3::from_array(t),
        glam::Quat::from_array(r),
        Vec3::from_array(s),
    )
}

/// Reads one primitive's vertices and indices into a [`Mesh`].
fn read_primitive(
    primitive: &gltf::Primitive,
    buffers: &[gltf::buffer::Data],
) -> Result<Mesh, GltfError> {
    let reader = primitive.reader(|buffer| Some(&buffers[buffer.index()]));

    let positions: Vec<[f32; 3]> = reader
        .read_positions()
        .ok_or(GltfError::MissingPositions)?
        .collect();

    let normals: Option<Vec<[f32; 3]>> = reader.read_normals().map(|n| n.collect());
    let uvs: Option<Vec<[f32; 2]>> = reader.read_tex_coords(0).map(|tc| tc.into_f32().collect());

    let indices: Vec<u32> = match reader.read_indices() {
        Some(idx) => idx.into_u32().collect(),
        // Non-indexed primitive: every three positions form a triangle.
        None => (0..positions.len() as u32).collect(),
    };

    // Generate flat normals when the document omits them, so lighting still works.
    let normals = normals.unwrap_or_else(|| compute_normals(&positions, &indices));

    let vertices = positions
        .iter()
        .enumerate()
        .map(|(i, &p)| {
            let n = normals.get(i).copied().unwrap_or([0.0, 1.0, 0.0]);
            let uv = uvs
                .as_ref()
                .and_then(|u| u.get(i))
                .copied()
                .unwrap_or([0.0, 0.0]);
            Vertex::new(
                Vec3::from_array(p),
                Vec3::from_array(n),
                Vec2::from_array(uv),
            )
        })
        .collect();

    Ok(Mesh::new(vertices, indices))
}

/// Smooth per-vertex normals: accumulate each triangle's face normal at its
/// vertices, then normalize.
fn compute_normals(positions: &[[f32; 3]], indices: &[u32]) -> Vec<[f32; 3]> {
    let mut normals = vec![Vec3::ZERO; positions.len()];
    for tri in indices.chunks_exact(3) {
        let (a, b, c) = (tri[0] as usize, tri[1] as usize, tri[2] as usize);
        let pa = Vec3::from_array(positions[a]);
        let pb = Vec3::from_array(positions[b]);
        let pc = Vec3::from_array(positions[c]);
        let face = (pb - pa).cross(pc - pa);
        normals[a] += face;
        normals[b] += face;
        normals[c] += face;
    }
    normals
        .into_iter()
        .map(|n| n.normalize_or_zero().to_array())
        .collect()
}

/// Maps a primitive's PBR metallic-roughness factors onto a [`Material`].
fn read_material(primitive: &gltf::Primitive) -> Material {
    let pbr = primitive.material().pbr_metallic_roughness();
    let [r, g, b, a] = pbr.base_color_factor();
    Material {
        albedo: Color::rgba(r, g, b, a),
        metallic: pbr.metallic_factor(),
        roughness: pbr.roughness_factor(),
    }
}