engine/asset/src/mesh.rs

use utilities::prelude::*;
use vulkan_rs::prelude::*;

use anyhow::Result;
use gltf_loader::prelude::*;
use utilities::prelude::cgmath::{prelude::One, Matrix4};

use std::sync::Arc;

use crate::scene_type::SceneType;
use crate::vertextypes::prelude::*;

use super::asset::MatrixBuffer;

pub const BONE_COUNT: usize = 64;

fn new_model_buffer(
    device: &Arc<Device>,
    model_matrix: impl Into<MatrixBuffer>,
) -> Result<Arc<Buffer<MatrixBuffer>>> {
    Buffer::builder()
        .set_usage(VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT)
        .set_memory_usage(MemoryUsage::CpuOnly)
        .set_data(&[model_matrix.into()])
        .build(device.clone())
}

fn new_bone_buffer(device: &Arc<Device>) -> Result<Arc<Buffer<MatrixBuffer>>> {
    Buffer::builder()
        .set_usage(VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT)
        .set_memory_usage(MemoryUsage::CpuOnly)
        .set_data(&[cgmath::Matrix4::one().into(); BONE_COUNT])
        .build(device.clone())
}

pub trait AddPrimitiveToMesh<T: ReprC + Copy + Send + Sync + 'static> {
    fn add_primitive(
        &mut self,
        vertex_buffer: Arc<Buffer<T>>,
        color_texture: Option<Arc<Image>>,
        normal_texture: Option<Arc<Image>>,
        material: PrimitiveMaterial,
        opaque: bool,
    ) -> Result<()>;
}

#[derive(Clone, Debug)]
pub struct AssetMesh {
    device: Arc<Device>,
    render_type: SceneType,

    pub name: Option<String>,

    pub node_index: usize,

    // model_matrix with uniform buffer
    model_matrix: Matrix4<f32>,
    pub(crate) model_buffer: Arc<Buffer<MatrixBuffer>>,

    pub(crate) bone_buffer: Option<Arc<Buffer<MatrixBuffer>>>,

    // primitives
    pub primitive_types: Vec<PrimitiveDataTypes>,
}

impl AssetMesh {
    pub fn new(device: &Arc<Device>, render_type: SceneType) -> Result<Self> {
        Self::new_custom_model(device, render_type, Matrix4::identity())
    }

    pub(crate) fn new_custom_model(
        device: &Arc<Device>,
        render_type: SceneType,
        model_matrix: Matrix4<f32>,
    ) -> Result<Self> {
        Ok(AssetMesh {
            device: device.clone(),
            render_type,

            name: None,

            node_index: 0,

            model_matrix,
            model_buffer: new_model_buffer(device, model_matrix)?,

            bone_buffer: None,

            primitive_types: Vec::new(),
        })
    }

    pub fn set_node_index(&mut self, node_index: usize) {
        self.node_index = node_index;
    }

    pub fn add_bone_buffer(&mut self) -> Result<()> {
        self.bone_buffer = Some(new_bone_buffer(&self.device)?);

        Ok(())
    }

    pub fn set_name(&mut self, name: &str) {
        self.name = Some(name.to_string())
    }

    pub fn bone_buffer(&self) -> &Option<Arc<Buffer<MatrixBuffer>>> {
        &self.bone_buffer
    }

    pub fn model_matrix(&self) -> Matrix4<f32> {
        self.model_matrix
    }

    pub fn model_buffer(&self) -> &Arc<Buffer<MatrixBuffer>> {
        &self.model_buffer
    }

    /// basically a clone, but it creates new vulkan handles to separate instances
    pub fn separated_copy(&self) -> Result<AssetMesh> {
        let model_buffer = new_model_buffer(&self.device, self.model_matrix)?;

        let bone_buffer = match self.bone_buffer {
            Some(_) => Some(new_bone_buffer(&self.device)?),
            None => None,
        };

        let mut primitive_types = Vec::with_capacity(self.primitive_types.len());

        for primitive_type in &self.primitive_types {
            primitive_types.push(primitive_type.clone(&model_buffer, bone_buffer.as_ref())?);
        }

        Ok(AssetMesh {
            device: self.device.clone(),
            render_type: self.render_type,

            name: self.name.clone(),

            node_index: self.node_index,

            model_matrix: self.model_matrix,
            model_buffer,

            bone_buffer,

            primitive_types,
        })
    }
}

impl AddPrimitiveToMesh<PositionOnly> for AssetMesh {
    fn add_primitive(
        &mut self,
        vertex_buffer: Arc<Buffer<PositionOnly>>,
        color_texture: Option<Arc<Image>>,
        normal_texture: Option<Arc<Image>>,
        material: PrimitiveMaterial,
        opaque: bool,
    ) -> Result<()> {
        self.primitive_types.push(PrimitiveDataTypes::new(
            &self.device,
            vertex_buffer,
            color_texture,
            normal_texture,
            &self.model_buffer,
            self.bone_buffer.as_ref(),
            self.render_type,
            material,
            opaque,
        )?);

        Ok(())
    }
}

impl AddPrimitiveToMesh<PositionNormal> for AssetMesh {
    fn add_primitive(
        &mut self,
        vertex_buffer: Arc<Buffer<PositionNormal>>,
        color_texture: Option<Arc<Image>>,
        normal_texture: Option<Arc<Image>>,
        material: PrimitiveMaterial,
        opaque: bool,
    ) -> Result<()> {
        self.primitive_types.push(PrimitiveDataTypes::new(
            &self.device,
            vertex_buffer,
            color_texture,
            normal_texture,
            &self.model_buffer,
            self.bone_buffer.as_ref(),
            self.render_type,
            material,
            opaque,
        )?);

        Ok(())
    }
}

impl AddPrimitiveToMesh<PositionNormalBone> for AssetMesh {
    fn add_primitive(
        &mut self,
        vertex_buffer: Arc<Buffer<PositionNormalBone>>,
        color_texture: Option<Arc<Image>>,
        normal_texture: Option<Arc<Image>>,
        material: PrimitiveMaterial,
        opaque: bool,
    ) -> Result<()> {
        self.primitive_types.push(PrimitiveDataTypes::new(
            &self.device,
            vertex_buffer,
            color_texture,
            normal_texture,
            &self.model_buffer,
            self.bone_buffer.as_ref(),
            self.render_type,
            material,
            opaque,
        )?);

        Ok(())
    }
}

impl AddPrimitiveToMesh<PositionNormalColor> for AssetMesh {
    fn add_primitive(
        &mut self,
        vertex_buffer: Arc<Buffer<PositionNormalColor>>,
        color_texture: Option<Arc<Image>>,
        normal_texture: Option<Arc<Image>>,
        material: PrimitiveMaterial,
        opaque: bool,
    ) -> Result<()> {
        self.primitive_types.push(PrimitiveDataTypes::new(
            &self.device,
            vertex_buffer,
            color_texture,
            normal_texture,
            &self.model_buffer,
            self.bone_buffer.as_ref(),
            self.render_type,
            material,
            opaque,
        )?);

        Ok(())
    }
}

impl AddPrimitiveToMesh<PositionColorUV> for AssetMesh {
    fn add_primitive(
        &mut self,
        vertex_buffer: Arc<Buffer<PositionColorUV>>,
        color_texture: Option<Arc<Image>>,
        normal_texture: Option<Arc<Image>>,
        material: PrimitiveMaterial,
        opaque: bool,
    ) -> Result<()> {
        self.primitive_types.push(PrimitiveDataTypes::new(
            &self.device,
            vertex_buffer,
            color_texture,
            normal_texture,
            &self.model_buffer,
            self.bone_buffer.as_ref(),
            self.render_type,
            material,
            opaque,
        )?);

        Ok(())
    }
}

impl AddPrimitiveToMesh<PositionNormalColorUV> for AssetMesh {
    fn add_primitive(
        &mut self,
        vertex_buffer: Arc<Buffer<PositionNormalColorUV>>,
        color_texture: Option<Arc<Image>>,
        normal_texture: Option<Arc<Image>>,
        material: PrimitiveMaterial,
        opaque: bool,
    ) -> Result<()> {
        self.primitive_types.push(PrimitiveDataTypes::new(
            &self.device,
            vertex_buffer,
            color_texture,
            normal_texture,
            &self.model_buffer,
            self.bone_buffer.as_ref(),
            self.render_type,
            material,
            opaque,
        )?);

        Ok(())
    }
}

impl AddPrimitiveToMesh<PositionNormalColorUVNormalUV> for AssetMesh {
    fn add_primitive(
        &mut self,
        vertex_buffer: Arc<Buffer<PositionNormalColorUVNormalUV>>,
        color_texture: Option<Arc<Image>>,
        normal_texture: Option<Arc<Image>>,
        material: PrimitiveMaterial,
        opaque: bool,
    ) -> Result<()> {
        self.primitive_types.push(PrimitiveDataTypes::new(
            &self.device,
            vertex_buffer,
            color_texture,
            normal_texture,
            &self.model_buffer,
            self.bone_buffer.as_ref(),
            self.render_type,
            material,
            opaque,
        )?);

        Ok(())
    }
}

impl AddPrimitiveToMesh<PositionNormalColorUVBone> for AssetMesh {
    fn add_primitive(
        &mut self,
        vertex_buffer: Arc<Buffer<PositionNormalColorUVBone>>,
        color_texture: Option<Arc<Image>>,
        normal_texture: Option<Arc<Image>>,
        material: PrimitiveMaterial,
        opaque: bool,
    ) -> Result<()> {
        self.primitive_types.push(PrimitiveDataTypes::new(
            &self.device,
            vertex_buffer,
            color_texture,
            normal_texture,
            &self.model_buffer,
            self.bone_buffer.as_ref(),
            self.render_type,
            material,
            opaque,
        )?);

        Ok(())
    }
}

impl AddPrimitiveToMesh<PositionNormalColorUVNormalUVBone> for AssetMesh {
    fn add_primitive(
        &mut self,
        vertex_buffer: Arc<Buffer<PositionNormalColorUVNormalUVBone>>,
        color_texture: Option<Arc<Image>>,
        normal_texture: Option<Arc<Image>>,
        material: PrimitiveMaterial,
        opaque: bool,
    ) -> Result<()> {
        self.primitive_types.push(PrimitiveDataTypes::new(
            &self.device,
            vertex_buffer,
            color_texture,
            normal_texture,
            &self.model_buffer,
            self.bone_buffer.as_ref(),
            self.render_type,
            material,
            opaque,
        )?);

        Ok(())
    }
}