engine/gltf-loader/src/primitive.rs

use super::skin::Skin;

use utilities::prelude::cgmath::{Matrix4, Vector2, Vector3, Vector4};

use gltf;
use gltf::{
    buffer::Buffer,
    material::{AlphaMode, NormalTexture},
    mesh::{util::ReadTexCoords, BoundingBox, Reader},
    texture::Info,
};

#[derive(Debug)]
pub struct TextureInfo {
    pub image_index: usize,
    pub texture_index: u32,
    pub texture_coordinates: Vec<Vector2<f32>>,
}

#[derive(Debug, Copy, Clone)]
pub struct Material {
    pub color: [f32; 4],
    pub metallic_factor: f32,
    pub emissive_factor: [f32; 3],
    pub roughness_factor: f32,
    pub alpha_mode: AlphaMode,
    pub alpha_cut_off: f32,
}

impl Default for Material {
    fn default() -> Self {
        Self {
            color: [1.0, 1.0, 1.0, 1.0],
            metallic_factor: 1.0,
            emissive_factor: [0.0, 0.0, 0.0],
            roughness_factor: 1.0,
            alpha_mode: AlphaMode::Blend,
            alpha_cut_off: 1.0,
        }
    }
}

#[derive(Debug)]
pub struct Primitive {
    pub positions: Option<Vec<Vector3<f32>>>,
    pub normals: Option<Vec<Vector3<f32>>>,
    pub tangents: Option<Vec<Vector4<f32>>>,
    pub indices: Option<Vec<u32>>,

    pub bounding_box: BoundingBox,

    pub color_texture_info: Option<TextureInfo>,
    pub normal_texture_info: Option<TextureInfo>,
    pub emissive_texture_info: Option<TextureInfo>,
    pub roughness_texture_info: Option<TextureInfo>,

    pub joint_indices: Option<Vec<[i32; 4]>>,
    pub weights: Option<Vec<[f32; 4]>>,

    pub material: Material,
}

impl Primitive {
    pub fn new(
        gltf_primitive: gltf::Primitive<'_>,
        buffers: &Vec<gltf::buffer::Data>,
        skin: Option<&Skin>,
    ) -> Primitive {
        // create reader
        let reader = gltf_primitive.reader(|buffer| Some(&buffers[buffer.index()]));

        // read positions
        let positions = match reader.read_positions() {
            Some(iter) => Some(iter.map(|vertex| Vector3::from(vertex)).collect()),
            None => None,
        };

        // read normals
        let normals = match reader.read_normals() {
            Some(iter) => Some(iter.map(|normal| Vector3::from(normal)).collect()),
            None => None,
        };

        // read tangents
        let tangents = match reader.read_tangents() {
            Some(iter) => Some(iter.map(|tangent| Vector4::from(tangent)).collect()),
            None => None,
        };

        // read indices
        let indices = match reader.read_indices() {
            Some(iter) => Some(iter.into_u32().map(|index| index).collect()),
            None => None,
        };

        let material = gltf_primitive.material();

        let normal_texture_info =
            Self::load_normal_texture_info(material.normal_texture(), &reader);

        let metallic_roughness = material.pbr_metallic_roughness();

        let color_texture_info =
            Self::load_info_texture_info(metallic_roughness.base_color_texture(), &reader);

        let emissive_texture_info =
            Self::load_info_texture_info(material.emissive_texture(), &reader);

        let roughness_texture_info =
            Self::load_info_texture_info(metallic_roughness.metallic_roughness_texture(), &reader);

        let material = Material {
            color: metallic_roughness.base_color_factor(),
            metallic_factor: metallic_roughness.metallic_factor(),
            emissive_factor: material.emissive_factor(),
            roughness_factor: metallic_roughness.roughness_factor(),
            alpha_mode: material.alpha_mode(),
            alpha_cut_off: material.alpha_cutoff().unwrap_or(1.0),
        };

        let (joint_indices, weights): (Option<Vec<[i32; 4]>>, Option<Vec<[f32; 4]>>) = match skin {
            Some(skin) => {
                let joint_ids = skin.joints();

                let joint_indices = match reader.read_joints(skin.index as u32) {
                    Some(iter) => Some(
                        iter.into_u16()
                            .map(|joint| {
                                let mut joints = [0, 0, 0, 0];

                                for j in 0..4 {
                                    joints[j] = joint_ids[joint[j] as usize] as i32;
                                }

                                joints
                            })
                            .collect(),
                    ),
                    None => None,
                };

                let weights = match reader.read_weights(skin.index as u32) {
                    Some(iter) => Some(
                        iter.into_f32()
                            .map(|weight| {
                                let mut weights = [0.0, 0.0, 0.0, 0.0];

                                let mut sum = 0.0;
                                for w in &weight {
                                    sum += w;
                                }
                                let inv_sum = 1.0 / sum;

                                for (j, w) in weight.iter().enumerate() {
                                    weights[j] = w * inv_sum;
                                }

                                weights
                            })
                            .collect(),
                    ),
                    None => None,
                };

                (joint_indices, weights)
            }
            None => (None, None),
        };

        if let Some(ref joint_indices) = joint_indices {
            if let Some(ref weights) = weights {
                debug_assert!(joint_indices.len() == weights.len());
            }
        }

        Primitive {
            positions,
            normals,
            tangents,
            indices,

            bounding_box: gltf_primitive.bounding_box(),

            color_texture_info,
            normal_texture_info,
            emissive_texture_info,
            roughness_texture_info,

            joint_indices,
            weights,

            material,
        }
    }

    pub(crate) fn scale_bb(&mut self, m: Matrix4<f32>) {
        self.bounding_box = BoundingBox {
            min: Self::mul(self.bounding_box.min, m),
            max: Self::mul(self.bounding_box.max, m),
        };
    }

    fn mul(v: impl Into<Vector3<f32>>, m: Matrix4<f32>) -> [f32; 3] {
        let t = m * v.into().extend(1.0);

        [t.x / t.w, t.y / t.w, t.z / t.w]
    }

    #[inline]
    fn load_info_texture_info<'a, 's, F>(
        info: Option<Info<'_>>,
        reader: &Reader<'a, 's, F>,
    ) -> Option<TextureInfo>
    where
        F: Clone + Fn(Buffer<'a>) -> Option<&'s [u8]>,
    {
        info.map(|info| {
            Self::load_texture_info(
                reader.read_tex_coords(info.tex_coord()),
                info.texture().source().index(),
                info.tex_coord(),
            )
        })
        .flatten()
    }

    #[inline]
    fn load_normal_texture_info<'a, 's, F>(
        info: Option<NormalTexture<'_>>,
        reader: &Reader<'a, 's, F>,
    ) -> Option<TextureInfo>
    where
        F: Clone + Fn(Buffer<'a>) -> Option<&'s [u8]>,
    {
        info.map(|info| {
            Self::load_texture_info(
                reader.read_tex_coords(info.tex_coord()),
                info.texture().source().index(),
                info.tex_coord(),
            )
        })
        .flatten()
    }

    #[inline]
    fn load_texture_info<'s>(
        read: Option<ReadTexCoords<'s>>,
        image_index: usize,
        texture_index: u32,
    ) -> Option<TextureInfo> {
        match read {
            Some(iter) => {
                let texture_coordinates = iter
                    .into_f32()
                    .map(|texture_coordinate| Vector2::from(texture_coordinate))
                    .collect();

                Some(TextureInfo {
                    image_index,
                    texture_index,
                    texture_coordinates,
                })
            }
            None => None,
        }
    }
}

impl PartialEq for Primitive {
    fn eq(&self, other: &Primitive) -> bool {
        self.positions.is_some() == other.positions.is_some()
            && self.normals.is_some() == other.normals.is_some()
            && self.indices.is_some() == other.indices.is_some()
            && self.color_texture_info.is_some() == other.color_texture_info.is_some()
            && self.normal_texture_info.is_some() == other.normal_texture_info.is_some()
            && self.joint_indices.is_some() == other.joint_indices.is_some()
            && self.weights.is_some() == other.weights.is_some()
    }
}