awsm-renderer 0.1.7

//! glTF scene population into renderer resources.

use std::{
    collections::{HashMap, HashSet},
    sync::{Arc, Mutex},
};

use awsm_renderer_core::texture::texture_pool::TextureColorInfo;
use glam::Mat4;

use crate::materials::MaterialKey;
use crate::{meshes::MeshKey, textures::TextureKey, transforms::TransformKey, AwsmRenderer};

use super::{data::GltfData, error::AwsmGltfError};

mod animation;
mod extensions;
pub mod material;
mod mesh;
mod skin;
pub(super) mod transforms;

/// Context and shared state used while populating glTF data.
pub struct GltfPopulateContext {
    pub data: Arc<GltfData>,
    pub textures: Mutex<HashMap<GltfTextureKey, TextureKey>>,
    pub(super) material_keys: Mutex<HashMap<GltfMaterialLookupKey, MaterialKey>>,
    pub node_to_skin_transform:
        Mutex<HashMap<GltfIndex, Arc<(Vec<TransformKey>, Vec<SkinInverseBindMatrix>)>>>,
    pub transform_is_joint: Mutex<HashSet<TransformKey>>,
    pub transform_is_instanced: Mutex<HashSet<TransformKey>>,
    pub(super) node_animation_samplers: HashMap<GltfIndex, GltfNodeAnimationSamplers>,
    pub key_lookups: Arc<Mutex<GltfKeyLookups>>,
}

/// Lookup tables for glTF node, mesh, and primitive keys.
#[derive(Debug, Clone, Default)]
pub struct GltfKeyLookups {
    pub node_transforms: HashMap<String, TransformKey>,
    // for all nodes with a name, get mesh_keys per primitive for that node, and optional mesh name
    pub node_meshes: HashMap<String, Vec<(Option<String>, Vec<MeshKey>)>>,
    // for all the meshes with a name, get mesh_keys per primitive for that mesh
    pub mesh_primitives: HashMap<String, Vec<MeshKey>>,
    pub node_index_to_transform: HashMap<GltfIndex, TransformKey>,
    pub all_mesh_keys: HashMap<GltfIndex, Vec<MeshKey>>,
}

impl GltfKeyLookups {
    /// Records a transform key for a glTF node.
    pub fn insert_transform(&mut self, node: &gltf::Node, key: TransformKey) {
        if let Some(name) = node.name() {
            self.node_transforms.insert(name.to_string(), key);
        }

        self.node_index_to_transform.insert(node.index(), key);
    }

    /// Records a mesh key for a glTF node and mesh.
    pub fn insert_mesh(&mut self, node: &gltf::Node, mesh: &gltf::Mesh, mesh_key: MeshKey) {
        self.all_mesh_keys
            .entry(mesh.index())
            .or_default()
            .push(mesh_key);

        if let Some(mesh_name) = mesh.name() {
            self.mesh_primitives
                .entry(mesh_name.to_string())
                .or_default()
                .push(mesh_key);
        }

        if let Some(node_name) = node.name() {
            let entry = self.node_meshes.entry(node_name.to_string()).or_default();
            match mesh.name() {
                None => {
                    // no mesh name, just add to the list with None
                    entry.push((None, vec![mesh_key]));
                }
                Some(name) => {
                    // see if we already have an entry for this mesh name
                    let mut found = false;
                    for (mesh_name_opt, mesh_keys) in entry.iter_mut() {
                        if let Some(mesh_name) = mesh_name_opt {
                            if mesh_name == name {
                                mesh_keys.push(mesh_key);
                                found = true;
                            }
                        }
                    }

                    // otherwise add a new entry
                    if !found {
                        entry.push((Some(name.to_string()), vec![mesh_key]));
                    }
                }
            }
        }
    }

    /// Returns an iterator over meshes for a node name.
    pub fn meshes_for_node_iter(&self, node_name: &str) -> impl Iterator<Item = &MeshKey> {
        self.node_meshes
            .get(node_name)
            .into_iter()
            .flat_map(|entries| entries.iter())
            .flat_map(|(_mesh_name_opt, mesh_keys)| mesh_keys.iter())
    }
}

/// Key that identifies a glTF texture plus color info.
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct GltfTextureKey {
    pub index: GltfIndex,
    pub color: TextureColorInfo,
}

type SkinInverseBindMatrix = Mat4;

type GltfIndex = usize;

#[derive(Clone, Copy, Debug)]
pub(super) struct GltfAnimationSamplerRef {
    pub animation_index: usize,
    pub channel_index: usize,
    pub sampler_index: usize,
}

#[derive(Clone, Debug, Default)]
pub(super) struct GltfNodeAnimationSamplers {
    pub translation: Option<GltfAnimationSamplerRef>,
    pub rotation: Option<GltfAnimationSamplerRef>,
    pub scale: Option<GltfAnimationSamplerRef>,
    pub morph: Option<GltfAnimationSamplerRef>,
}

#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
pub(super) struct GltfMaterialLookupKey {
    pub material_index: Option<usize>,
    pub vertex_color_set_index: Option<usize>,
    pub hud: bool,
}

impl AwsmRenderer {
    /// Populates renderer resources from a glTF asset.
    pub async fn populate_gltf(
        &mut self,
        gltf_data: impl Into<Arc<GltfData>>,
        scene: Option<usize>,
    ) -> anyhow::Result<GltfPopulateContext> {
        let gltf_data = gltf_data.into();
        self.gltf.raw_datas.push(gltf_data.clone());

        let mut mesh_keys = Vec::new();
        let node_animation_samplers = build_node_animation_sampler_lookup(&gltf_data.doc);

        let ctx = GltfPopulateContext {
            data: gltf_data,
            textures: Mutex::new(HashMap::new()),
            material_keys: Mutex::new(HashMap::new()),
            node_to_skin_transform: Mutex::new(HashMap::new()),
            transform_is_joint: Mutex::new(HashSet::new()),
            transform_is_instanced: Mutex::new(HashSet::new()),
            node_animation_samplers,
            key_lookups: Arc::new(Mutex::new(GltfKeyLookups::default())),
        };

        let scene = match scene {
            Some(index) => ctx
                .data
                .doc
                .scenes()
                .nth(index)
                .ok_or(AwsmGltfError::InvalidScene(index))?,
            None => match ctx.data.doc.default_scene() {
                Some(scene) => scene,
                None => ctx
                    .data
                    .doc
                    .scenes()
                    .next()
                    .ok_or(AwsmGltfError::NoDefaultScene)?,
            },
        };

        for node in scene.nodes() {
            self.populate_gltf_node_transform(&ctx, &node, None)?;
        }

        for node in scene.nodes() {
            self.populate_gltf_node_extension_instancing(&ctx, &node)?;
        }

        for node in scene.nodes() {
            self.populate_gltf_node_skin(&ctx, &node)?;
        }

        for node in scene.nodes() {
            self.populate_gltf_node_animation(&ctx, &node)?;
        }

        for node in scene.nodes() {
            mesh_keys.push(self.populate_gltf_node_mesh(&ctx, &node).await?);
        }

        self.finalize_gpu_textures().await?;

        Ok(ctx)
    }
}

impl GltfPopulateContext {
    pub(super) fn resolve_animation_sampler(
        &self,
        sampler_ref: GltfAnimationSamplerRef,
    ) -> Result<gltf::animation::Sampler<'_>, AwsmGltfError> {
        self.data
            .doc
            .animations()
            .nth(sampler_ref.animation_index)
            .and_then(|animation| animation.samplers().nth(sampler_ref.sampler_index))
            .ok_or(AwsmGltfError::MissingAnimationSampler {
                animation_index: sampler_ref.animation_index,
                channel_index: sampler_ref.channel_index,
                sampler_index: sampler_ref.sampler_index,
            })
    }
}

fn build_node_animation_sampler_lookup(
    doc: &gltf::Document,
) -> HashMap<GltfIndex, GltfNodeAnimationSamplers> {
    let mut out = HashMap::<GltfIndex, GltfNodeAnimationSamplers>::new();

    for animation in doc.animations() {
        for channel in animation.channels() {
            let node_index = channel.target().node().index();
            let entry = out.entry(node_index).or_default();
            let sampler_ref = GltfAnimationSamplerRef {
                animation_index: animation.index(),
                channel_index: channel.index(),
                sampler_index: channel.sampler().index(),
            };

            match channel.target().property() {
                gltf::animation::Property::Translation => {
                    if entry.translation.is_none() {
                        entry.translation = Some(sampler_ref);
                    }
                }
                gltf::animation::Property::Rotation => {
                    if entry.rotation.is_none() {
                        entry.rotation = Some(sampler_ref);
                    }
                }
                gltf::animation::Property::Scale => {
                    if entry.scale.is_none() {
                        entry.scale = Some(sampler_ref);
                    }
                }
                gltf::animation::Property::MorphTargetWeights => {
                    if entry.morph.is_none() {
                        entry.morph = Some(sampler_ref);
                    }
                }
            }
        }
    }

    out
}