path: root/src/classes/mesh.rs

use super::streaming_info::StreamingInfo;
use crate::object::{Value, parser::FromValue};
use anyhow::{Result, anyhow, bail};
use glam::{Mat4, Vec2, Vec3, Vec3A, Vec4};
use log::debug;
use serde::Serialize;
use std::mem::transmute;

#[derive(Debug, Serialize)]
pub struct Mesh {
    pub name: String,
    pub bind_pose: Vec<Mat4>,
    pub bone_name_hashes: Vec<u32>,
    pub index_format: i32,
    pub index_buffer: Vec<u8>,
    pub sub_meshes: Vec<SubMesh>,
    pub stream_data: StreamingInfo,
    pub vertex_data: VertexData,
}

#[derive(Debug, Serialize)]
pub struct SubMesh {
    pub topology: i32,
    pub vertex_count: u32,
    pub base_vertex: u32,
    pub first_byte: u32,
    pub first_vertex: u32,
    pub index_count: u32,
}

#[derive(Debug, Serialize)]
pub struct VertexData {
    pub channels: Vec<ChannelInfo>,
    pub data: Vec<u8>,
    pub vertex_count: u32,
}

#[derive(Debug, Serialize)]
pub struct ChannelInfo {
    pub dimension: u8,
    pub format: VertexFormat,
    pub offset: u8,
    pub stream: u8,
}

impl FromValue for Mesh {
    fn from_value(v: Value) -> Result<Self> {
        let mut fields = v.as_class("Mesh")?;
        Ok(Mesh {
            name: fields.field("m_Name")?,
            index_format: fields.field("m_IndexFormat")?,
            vertex_data: fields.field("m_VertexData")?,
            stream_data: fields.field("m_StreamData")?,
            index_buffer: fields
                .remove("m_IndexBuffer")
                .unwrap()
                .as_vector()
                .unwrap()
                .into_iter()
                .map(|e| e.as_u8().unwrap())
                .collect(),
            sub_meshes: fields
                .remove("m_SubMeshes")
                .unwrap()
                .as_vector()
                .unwrap()
                .into_iter()
                .map(|e| e.parse().unwrap())
                .collect(),
            bind_pose: fields
                .remove("m_BindPose")
                .unwrap()
                .as_vector()
                .unwrap()
                .into_iter()
                .map(|e| e.parse().unwrap())
                .collect(),
            bone_name_hashes: fields
                .remove("m_BoneNameHashes")
                .unwrap()
                .as_vector()
                .unwrap()
                .into_iter()
                .map(|e| e.parse().unwrap())
                .collect(),
        })
    }
}

impl Mesh {
    pub fn read_indecies(&self) -> Vec<[u32; 3]> {
        if self.index_format == 0 {
            self.index_buffer
                .array_chunks::<2>()
                .map(|x| u16::from_le_bytes(*x) as u32)
                .array_chunks()
                .collect()
        } else {
            self.index_buffer
                .array_chunks::<4>()
                .map(|x| u32::from_le_bytes(*x))
                .array_chunks()
                .collect()
        }
    }
}

impl FromValue for VertexData {
    fn from_value(v: Value) -> Result<Self> {
        let mut fields = v.as_class("VertexData")?;
        Ok(VertexData {
            vertex_count: fields.field("m_VertexCount")?,
            data: fields
                .remove("m_DataSize")
                .ok_or(anyhow!("m_DataSize missing"))?
                .as_typeless()
                .ok_or(anyhow!("m_DataSize is not typeless"))?,
            channels: fields
                .remove("m_Channels")
                .unwrap()
                .as_vector()
                .unwrap()
                .into_iter()
                .map(|e| e.parse().unwrap())
                .collect(),
        })
    }
}

impl VertexData {
    /// Returns (offset, stride) for each stream.
    pub fn stream_layout(&self) -> Vec<(usize, usize)> {
        let stream_count = self.channels.iter().map(|c| c.stream).max().unwrap() + 1;
        let mut streams = Vec::new();
        let mut offset = 0;
        for si in 0..stream_count {
            let stride = self
                .channels
                .iter()
                .filter(|c| c.stream == si)
                // The modulo operator here is a hack to fix normal with 52 dimensions to 4
                .map(|c| (c.dimension as usize % 48) * c.format.component_size())
                .sum();
            streams.push((offset, stride));
            offset += stride * self.vertex_count as usize
        }
        streams
    }
    /// Reads a vertex channel and returns dimension count and data converted to floats
    pub fn read_channel(&self, ch: VertexDataChannel) -> Option<(usize, Vec<f32>)> {
        let channel = &self.channels[ch as u8 as usize];
        if channel.dimension == 0 {
            return None;
        }
        let dim = channel.dimension as usize % 48;
        let component_offset = channel.offset as usize;
        let component_size = channel.format.component_size();
        let (offset, stride) = self.stream_layout()[channel.stream as usize];
        debug!(
            "reading {ch:?} vertex channel (stride={stride}, offset={offset}, format={:?})",
            channel.format
        );
        let mut out = Vec::new();
        for vi in 0..self.vertex_count as usize {
            for di in 0..dim {
                let off = offset + vi * stride + component_offset + component_size * di;
                let e = &self.data[off..];
                out.push(match channel.format {
                    VertexFormat::Float => f32::from_le_bytes([e[0], e[1], e[2], e[3]]),
                    VertexFormat::Float16 => f16::from_le_bytes([e[0], e[1]]) as f32,
                    VertexFormat::SInt8 => i8::from_le_bytes([e[0]]) as f32,
                    VertexFormat::UInt8 => u8::from_le_bytes([e[0]]) as f32,
                    VertexFormat::SInt16 => i16::from_le_bytes([e[0], e[1]]) as f32,
                    VertexFormat::UInt16 => u16::from_le_bytes([e[0], e[1]]) as f32,
                    VertexFormat::SInt32 => i32::from_le_bytes([e[0], e[1], e[2], e[3]]) as f32,
                    VertexFormat::UInt32 => u32::from_le_bytes([e[0], e[1], e[2], e[3]]) as f32,
                    VertexFormat::UNorm8 => u8::from_le_bytes([e[0]]) as f32 / 256.,
                    VertexFormat::SNorm8 => i8::from_le_bytes([e[0]]) as f32 / 128.,
                    VertexFormat::UNorm16 => u16::from_le_bytes([e[0], e[1]]) as f32 / 65536.,
                    VertexFormat::SNorm16 => i16::from_le_bytes([e[0], e[1]]) as f32 / 32768.,
                })
            }
        }
        Some((dim, out))
    }
    pub fn read_channel_vec<T: VectorType>(
        &self,
        channel: VertexDataChannel,
        coerce: bool,
    ) -> Result<Option<Vec<T>>> {
        let Some((dim, mut data)) = self.read_channel(channel) else {
            return Ok(None);
        };
        if dim != T::DIM {
            if coerce {
                let mut ndata = Vec::new();
                if dim > T::DIM {
                    for e in data.chunks_exact(dim) {
                        ndata.extend(&e[..T::DIM]);
                    }
                } else if dim < T::DIM {
                    for e in data.chunks_exact(dim) {
                        ndata.extend(e);
                        for _ in 0..(T::DIM - dim) {
                            ndata.push(0.);
                        }
                    }
                }
                data = ndata;
            } else {
                bail!(
                    "dimension mismatch reading {channel:?} channel ({} != {})",
                    dim,
                    T::DIM
                );
            }
        }
        Ok(Some(VectorType::convert_array(data)))
    }
}

pub trait VectorType: Sized {
    const DIM: usize;
    fn convert_array(a: Vec<f32>) -> Vec<Self>;
}
impl VectorType for Vec3A {
    const DIM: usize = 3;
    fn convert_array(a: Vec<f32>) -> Vec<Self> {
        a.into_iter()
            .array_chunks()
            .map(Vec3A::from_array)
            .collect()
    }
}
impl VectorType for Vec3 {
    const DIM: usize = 3;
    fn convert_array(a: Vec<f32>) -> Vec<Self> {
        a.into_iter().array_chunks().map(Vec3::from_array).collect()
    }
}
impl VectorType for Vec2 {
    const DIM: usize = 2;
    fn convert_array(a: Vec<f32>) -> Vec<Self> {
        a.into_iter().array_chunks().map(Vec2::from_array).collect()
    }
}
impl VectorType for Vec4 {
    const DIM: usize = 4;
    fn convert_array(a: Vec<f32>) -> Vec<Self> {
        a.into_iter().array_chunks().map(Vec4::from_array).collect()
    }
}

impl FromValue for ChannelInfo {
    fn from_value(v: Value) -> Result<Self> {
        let mut fields = v.as_class("ChannelInfo")?;
        Ok(ChannelInfo {
            dimension: fields.field("dimension")?,
            format: fields.field("format")?,
            offset: fields.field("offset")?,
            stream: fields.field("stream")?,
        })
    }
}

impl FromValue for SubMesh {
    fn from_value(v: Value) -> Result<Self> {
        let mut fields = v.as_class("SubMesh")?;
        Ok(SubMesh {
            topology: fields.field("topology")?,
            vertex_count: fields.field("vertexCount")?,
            base_vertex: fields.field("baseVertex")?,
            first_byte: fields.field("firstByte")?,
            first_vertex: fields.field("firstVertex")?,
            index_count: fields.field("indexCount")?,
        })
    }
}

#[repr(u8)]
#[derive(Debug, Serialize, Clone, Copy, PartialEq)]
pub enum VertexDataChannel {
    Position,
    Normal,
    Tangent,
    Color,
    TexCoord0,
    TexCoord1,
    TexCoord2,
    TexCoord3,
    TexCoord4,
    TexCoord5,
    TexCoord6,
    TexCoord7,
    BlendWeight,
    BlendIndices,
}

#[repr(u8)]
#[derive(Debug, Serialize, Clone, Copy, PartialEq)]
pub enum VertexFormat {
    Float,
    Float16,
    UNorm8,
    SNorm8,
    UNorm16,
    SNorm16,
    UInt8,
    SInt8,
    UInt16,
    SInt16,
    UInt32,
    SInt32,
}

impl FromValue for VertexFormat {
    fn from_value(v: Value) -> Result<Self> {
        let x = v.as_u8().ok_or(anyhow!("expected u8 vertex format"))?;
        if x < 12 {
            Ok(unsafe { transmute(x) })
        } else {
            bail!("unknown vertex format")
        }
    }
}

impl VertexFormat {
    pub fn component_size(&self) -> usize {
        use VertexFormat::*;
        match self {
            Float | UInt32 | SInt32 => 4,
            Float16 | UInt16 | SInt16 | UNorm16 | SNorm16 => 2,
            UInt8 | SInt8 | UNorm8 | SNorm8 => 1,
        }
    }
}