wavefront_loader 0.2.4

#![doc = include_str!("../README.md")]

use std::collections::HashMap;
use std::fmt::{Debug, Display};
use std::fs::File;
use std::io::Write;
use std::io::{BufRead, BufReader};
use std::ops::AddAssign;
use std::path::PathBuf;

use linear_isomorphic::{RefIterable, VectorSpace};
use mtl::MtlData;
use nalgebra;

pub mod mtl;

type Vec2 = nalgebra::Vector2<f32>;
type Vec3 = nalgebra::Vector3<f32>;

#[derive(Debug)]
pub struct ObjData
{
    pub vertices: Vec<Vec3>,
    pub point_colors: Vec<Vec3>,
    pub normals: Vec<Vec3>,
    pub uvs: Vec<Vec2>,
    pub vertex_face_indices: Vec<Vec<u64>>,
    pub normal_face_indices: Vec<Vec<u64>>,
    pub uv_face_indices: Vec<Vec<u64>>,
    pub lines: Vec<[u64; 2]>,
    pub objects: Vec<ObjectRanges>,
    pub face_material_use: Vec<usize>,
    pub material_list: Vec<MtlData>,
}

#[derive(Debug)]
pub struct ObjectRanges
{
    pub name: String,
    pub positions: (u64, u64),
    pub normals: (u64, u64),
    pub uvs: (u64, u64),
    pub lines: (u64, u64),
    pub polygons: [(u64, u64); 3],
}

impl ObjData
{
    pub fn from_disk_file(path: &str) -> Self
    {
        let file_path = PathBuf::from(path);
        let file = File::open(&file_path).expect(&format!("Cannot find {}.", path));
        let root = file_path.parent().unwrap();
        let reader = BufReader::new(file);

        let mut vert_topology = Vec::<Vec<u64>>::new();
        let mut normal_topology = Vec::<Vec<u64>>::new();
        let mut uv_topology = Vec::<Vec<u64>>::new();
        let mut lines = Vec::<[u64; 2]>::new();

        let mut vertices = Vec::<Vec3>::new();
        let mut colors = Vec::<Vec3>::new();
        let mut normals = Vec::<Vec3>::new();
        let mut uvs = Vec::<Vec2>::new();
        let mut object_ranges = Vec::<ObjectRanges>::new();
        let mut material_list = Vec::new();
        let mut material_index_map = HashMap::new();
        let mut active_material_index = usize::MAX;
        let mut face_material_use = Vec::new();
        for line in reader.lines()
        {
            let line_str = line.unwrap();
            let line_str = line_str.trim();
            let tokens = line_str.split_whitespace().collect::<Vec<&str>>();
            if tokens.is_empty()
            {
                continue;
            }
            match tokens[0]
            {
                "mtllib" =>
                {
                    let path = root.join(tokens[1]);
                    let mtl_list: Vec<_> =
                        MtlData::from_disk_file(path.as_os_str().to_str().unwrap())
                            .into_iter()
                            .filter(|mtl| !material_index_map.contains_key(&mtl.name))
                            .collect();
                    for (i, mtl) in mtl_list.iter().enumerate()
                    {
                        material_index_map
                            .insert(mtl.name.clone(), material_list.len() + i);
                    }
                    material_list.extend(mtl_list);
                }
                "usemtl" =>
                {
                    active_material_index =
                        *material_index_map.get(&tokens[1].to_string()).unwrap();
                }
                "o" =>
                {
                    object_ranges.push(ObjectRanges {
                        name: tokens[1..].join(" "),
                        positions: (vertices.len() as u64, 0),
                        normals: (normals.len() as u64, 0),
                        uvs: (uvs.len() as u64, 0),
                        lines: (lines.len() as u64, 0),
                        polygons: [
                            (vert_topology.len() as u64, 0),
                            (normal_topology.len() as u64, 0),
                            (uv_topology.len() as u64, 0),
                        ],
                    });
                }
                "v" =>
                {
                    add_vec_3d(&tokens[1..], &mut vertices);
                    if tokens[1..].len() == 6
                    {
                        add_vec_3d(&tokens[4..], &mut colors);
                    }
                }
                "vn" =>
                {
                    add_vec_3d(&tokens[1..], &mut normals);
                }
                "vt" =>
                {
                    add_vec_2d(&tokens[1..], &mut uvs);
                }
                "l" =>
                {
                    add_line(&tokens[1..], &mut lines);
                }
                "f" =>
                {
                    face_material_use.push(active_material_index);
                    add_face(
                        &tokens[1..],
                        &mut vert_topology,
                        &mut normal_topology,
                        &mut uv_topology,
                        vertices.len(),
                        uvs.len(),
                        normals.len(),
                    )
                }
                _ => continue,
            }

            if let Some(or) = object_ranges.last_mut()
            {
                or.positions.1 = vertices.len() as u64;
                or.normals.1 = normals.len() as u64;
                or.uvs.1 = uvs.len() as u64;
                or.lines.1 = lines.len() as u64;

                or.polygons[0].1 = vert_topology.len() as u64;
                or.polygons[1].1 = normal_topology.len() as u64;
                or.polygons[2].1 = uv_topology.len() as u64;
            }
        }

        ObjData {
            vertices,
            point_colors: colors,
            normals,
            uvs,
            vertex_face_indices: vert_topology,
            normal_face_indices: normal_topology,
            uv_face_indices: uv_topology,
            lines,
            objects: object_ranges,
            material_list,
            face_material_use,
        }
    }

    pub fn export<'a, T>(mesh: &'a T, path: &str)
    where
        T: WaveFrontCompatible<'a>,
    {
        std::fs::create_dir_all(std::path::Path::new(path).parent().unwrap()).unwrap();
        let mut file = File::create(path).unwrap();

        let one = 1;
        if mesh.point_color_iterator().count() == 0
        {
            for point in mesh.pos_iterator()
            {
                writeln!(file, "v {} {} {}", point[0], point[1], point[2]).unwrap();
            }
        }
        else
        {
            for (point, color) in mesh.pos_iterator().zip(mesh.point_color_iterator())
            {
                writeln!(
                    file,
                    "v {} {} {} {} {} {}",
                    point[0], point[1], point[2], color[0], color[1], color[2]
                )
                .unwrap();
            }
        }

        for norm in mesh.norm_iterator()
        {
            writeln!(file, "vn {} {} {}", norm[0], norm[1], norm[2]).unwrap();
        }

        for uv in mesh.uv_iterator()
        {
            writeln!(file, "vt {} {}", uv[0], uv[1]).unwrap();
        }

        for line in mesh.segment_iterator()
        {
            writeln!(file, "l {} {}", line[0] + one, line[1] + one).unwrap();
        }

        let mut indices = Vec::new();
        let mut vert_face_count = 0;
        for face in mesh.pos_index_iterator()
        {
            vert_face_count += 1;

            indices.push(Vec::new());
            let face_ids: &mut Vec<_> = indices.last_mut().unwrap();

            for pos_id in face
            {
                face_ids.push([Some(pos_id), None, None]);
            }
        }

        let mut uv_face_count = 0;
        for face in mesh.uv_index_iterator()
        {
            for (local_id, uv_id) in face.enumerate()
            {
                indices[uv_face_count][local_id][1] = Some(uv_id);
            }
            uv_face_count += 1;
        }

        debug_assert!(uv_face_count == 0 || uv_face_count == vert_face_count);

        let mut norm_face_count = 0;
        for face in mesh.norm_index_iterator()
        {
            for (local_id, norm_id) in face.enumerate()
            {
                indices[norm_face_count][local_id][2] = Some(norm_id);
            }
            norm_face_count += 1;
        }

        debug_assert!(norm_face_count == 0 || norm_face_count == vert_face_count);

        for face in indices
        {
            let mut face_data = "f ".to_string();

            for [pos_id, uv_id, norm_id] in face
            {
                let pos_str = (pos_id.unwrap() + one).to_string();
                let uv_str = match uv_id
                {
                    None => "".to_string(),
                    Some(x) => (x + one).to_string(),
                };
                let norm_str = match norm_id
                {
                    None => "".to_string(),
                    Some(x) => (x + one).to_string(),
                };

                face_data
                    .push_str(format!("{}/{}/{} ", pos_str, uv_str, norm_str).as_str());
            }

            face_data.push('\n');

            write!(file, "{}", face_data).unwrap();
        }
    }
}

pub trait WaveFrontCompatible<'a>
{
    type Scalar: num_traits::Float + Debug + AddAssign + Display;

    fn pos_iterator(&'a self) -> impl Iterator<Item = [Self::Scalar; 3]>
    {
        std::iter::empty()
    }
    fn point_color_iterator(&'a self) -> impl Iterator<Item = [Self::Scalar; 3]>
    {
        std::iter::empty()
    }
    fn uv_iterator(&'a self) -> impl Iterator<Item = [Self::Scalar; 2]>
    {
        std::iter::empty()
    }
    fn norm_iterator(&'a self) -> impl Iterator<Item = [Self::Scalar; 3]>
    {
        std::iter::empty()
    }

    fn segment_iterator(&'a self) -> impl Iterator<Item = [usize; 2]>
    {
        std::iter::empty()
    }
    fn pos_index_iterator(&'a self) -> impl Iterator<Item = impl Iterator<Item = usize>>
    {
        let empty_iterator: std::iter::Empty<std::iter::Empty<usize>> =
            std::iter::empty();
        empty_iterator
    }
    fn uv_index_iterator(&'a self) -> impl Iterator<Item = impl Iterator<Item = usize>>
    {
        let empty_iterator: std::iter::Empty<std::iter::Empty<usize>> =
            std::iter::empty();
        empty_iterator
    }
    fn norm_index_iterator(&'a self)
    -> impl Iterator<Item = impl Iterator<Item = usize>>
    {
        let empty_iterator: std::iter::Empty<std::iter::Empty<usize>> =
            std::iter::empty();
        empty_iterator
    }
}

impl<'a> WaveFrontCompatible<'a> for ObjData
{
    type Scalar = f32;

    fn pos_iterator(&'a self) -> impl Iterator<Item = [f32; 3]>
    {
        self.vertices.iter().map(|v| [v.x, v.y, v.z])
    }

    fn point_color_iterator(&'a self) -> impl Iterator<Item = [f32; 3]>
    {
        self.point_colors.iter().map(|v| [v.x, v.y, v.z])
    }

    fn uv_iterator(&'a self) -> impl Iterator<Item = [f32; 2]>
    {
        self.uvs.iter().map(|v| [v.x, v.y])
    }

    fn norm_iterator(&'a self) -> impl Iterator<Item = [f32; 3]>
    {
        self.normals.iter().map(|v| [v.x, v.y, v.z])
    }

    fn segment_iterator(&'a self) -> impl Iterator<Item = [usize; 2]>
    {
        self.lines.iter().map(|[i, j]| [*i as usize, *j as usize])
    }

    fn pos_index_iterator(&'a self) -> impl Iterator<Item = impl Iterator<Item = usize>>
    {
        self.vertex_face_indices
            .iter()
            .map(|l| l.iter().map(|i| *i as usize))
    }

    fn uv_index_iterator(&'a self) -> impl Iterator<Item = impl Iterator<Item = usize>>
    {
        self.uv_face_indices
            .iter()
            .map(|l| l.iter().map(|i| *i as usize))
    }

    fn norm_index_iterator(&'a self)
    -> impl Iterator<Item = impl Iterator<Item = usize>>
    {
        self.normal_face_indices
            .iter()
            .map(|l| l.iter().map(|i| *i as usize))
    }
}

impl<'a, V, S> WaveFrontCompatible<'a> for Vec<V>
where
    V: VectorSpace<Scalar = S>,
    S: num_traits::Float + AddAssign + Display + Debug,
{
    type Scalar = S;

    fn pos_iterator(&'a self) -> impl Iterator<Item = [Self::Scalar; 3]>
    {
        self.iter().map(|v| [v[0], v[1], v[2]])
    }
}

impl<'a, V, S> WaveFrontCompatible<'a> for (&Vec<V>, &Vec<usize>)
where
    V: VectorSpace<Scalar = S>,
    S: num_traits::Float + AddAssign + Display + Debug,
{
    type Scalar = S;

    fn pos_iterator(&'a self) -> impl Iterator<Item = [Self::Scalar; 3]>
    {
        self.0.iter().map(|v| [v[0], v[1], v[2]])
    }

    fn pos_index_iterator(&'a self) -> impl Iterator<Item = impl Iterator<Item = usize>>
    {
        debug_assert!(self.1.len() % 3 == 0);
        self.1.chunks(3).map(|chunk| chunk.iter().copied())
    }
}

impl<'a, V, S, I> WaveFrontCompatible<'a> for (&Vec<V>, &Vec<Vec<I>>)
where
    V: VectorSpace<Scalar = S>,
    S: num_traits::Float + AddAssign + Display + Debug,
    usize: TryFrom<I>,
    I: num_traits::PrimInt + std::fmt::Display,
{
    type Scalar = S;

    fn pos_iterator(&'a self) -> impl Iterator<Item = [Self::Scalar; 3]>
    {
        self.0.iter().map(|v| [v[0], v[1], v[2]])
    }

    fn pos_index_iterator(&'a self) -> impl Iterator<Item = impl Iterator<Item = usize>>
    {
        self.1.iter().map(|face| {
            face.iter()
                .map(|i| usize::try_from(*i).unwrap_or(usize::MAX))
        })
    }
}

impl<'a, V, S, I> WaveFrontCompatible<'a> for (Vec<V>, Vec<[I; 3]>)
where
    V: VectorSpace<Scalar = S>,
    S: num_traits::Float + AddAssign + Display + Debug,
    usize: TryFrom<I>,
    I: num_traits::PrimInt + std::fmt::Display,
{
    type Scalar = S;

    fn pos_iterator(&'a self) -> impl Iterator<Item = [Self::Scalar; 3]>
    {
        self.0.iter().map(|v| [v[0], v[1], v[2]])
    }

    fn pos_index_iterator(&'a self) -> impl Iterator<Item = impl Iterator<Item = usize>>
    {
        self.1.iter().map(|face| {
            face.iter()
                .map(|i| usize::try_from(*i).unwrap_or(usize::MAX))
        })
    }
}


impl<'a, V, S, I> WaveFrontCompatible<'a> for (&Vec<V>, &Vec<[I; 3]>)
where
    V: VectorSpace<Scalar = S>,
    S: num_traits::Float + AddAssign + Display + Debug,
    usize: TryFrom<I>,
    I: num_traits::PrimInt + std::fmt::Display,
{
    type Scalar = S;

    fn pos_iterator(&'a self) -> impl Iterator<Item = [Self::Scalar; 3]>
    {
        self.0.iter().map(|v| [v[0], v[1], v[2]])
    }

    fn pos_index_iterator(&'a self) -> impl Iterator<Item = impl Iterator<Item = usize>>
    {
        self.1.iter().map(|face| {
            face.iter()
                .map(|i| usize::try_from(*i).unwrap_or(usize::MAX))
        })
    }
}

impl<'a, V, S, I> WaveFrontCompatible<'a> for (&Vec<V>, &Vec<[I; 2]>)
where
    V: VectorSpace<Scalar = S>,
    S: num_traits::Float + AddAssign + Display + Debug,
    usize: TryFrom<I>,
    I: num_traits::PrimInt + Display,
{
    type Scalar = S;

    fn pos_iterator(&'a self) -> impl Iterator<Item = [Self::Scalar; 3]>
    {
        self.0.iter().map(|v| [v[0], v[1], v[2]])
    }

    fn segment_iterator(&'a self) -> impl Iterator<Item = [usize; 2]>
    {
        self.1.iter().map(|[i, j]| {
            [
                usize::try_from(*i).unwrap_or(usize::MAX),
                usize::try_from(*j).unwrap_or(usize::MAX),
            ]
        })
    }
}

pub struct ColoredVerts<'a, V>
{
    pub points: &'a Vec<V>,
    pub colors: &'a Vec<V>,
}

impl<'a, V, S> WaveFrontCompatible<'a> for ColoredVerts<'a, V>
where
    V: VectorSpace<Scalar = S>,
    S: num_traits::Float + AddAssign + Display + Debug,
{
    type Scalar = S;

    fn pos_iterator(&'a self) -> impl Iterator<Item = [Self::Scalar; 3]>
    {
        self.points.iter().map(|v| [v[0], v[1], v[2]])
    }

    fn point_color_iterator(&'a self) -> impl Iterator<Item = [Self::Scalar; 3]>
    {
        self.colors.iter().map(|v| [v[0], v[1], v[2]])
    }
}

fn add_vec_3d(tokens: &[&str], vecs: &mut Vec<Vec3>)
{
    let vec = Vec3::new(
        tokens[0].parse::<f32>().unwrap(),
        tokens[1].parse::<f32>().unwrap(),
        tokens[2].parse::<f32>().unwrap(),
    );
    vecs.push(vec);
}

fn add_vec_2d(tokens: &[&str], vecs: &mut Vec<Vec2>)
{
    let vec = Vec2::new(
        tokens[0].parse::<f32>().unwrap(),
        tokens[1].parse::<f32>().unwrap(),
    );
    vecs.push(vec);
}

fn add_line(tokens: &[&str], vecs: &mut Vec<[u64; 2]>)
{
    assert!(tokens.len() == 2, "have {} expected 2", tokens.len());
    let index_1 = tokens[0].parse::<u64>().unwrap() - 1;
    let index_2 = tokens[1].parse::<u64>().unwrap() - 1;

    vecs.push([index_1, index_2])
}

fn add_face(
    tokens: &[&str],
    vert_topology: &mut Vec<Vec<u64>>,
    normal_topology: &mut Vec<Vec<u64>>,
    uv_topology: &mut Vec<Vec<u64>>,
    vert_count: usize,
    uv_count: usize,
    normal_count: usize,
)
{
    assert!(
        tokens.len() >= 3,
        "have {} out of a minimum of 3",
        tokens.len()
    );
    vert_topology.push(Vec::new());
    uv_topology.push(Vec::new());
    normal_topology.push(Vec::new());

    for token in tokens
    {
        let inner_tokens = token.split("/").collect::<Vec<&str>>();

        assert!(inner_tokens.len() <= 3 && inner_tokens.len() >= 1);

        let index = inner_tokens[0].parse::<i64>().unwrap();
        let vert_index = if index < 0
        {
            (vert_count as i64 + index) as u64
        }
        else
        {
            (index - 1) as u64
        };

        vert_topology.last_mut().unwrap().push(vert_index);

        // Only positions are specified.
        if inner_tokens.len() == 1
        {
            continue;
        }

        // Only positions and texture coordinates are specified.
        if inner_tokens.len() == 2
        {
            let index = inner_tokens[1].parse::<i64>().unwrap();
            let uv_index = if index < 0
            {
                (uv_count as i64 + index) as u64
            }
            else
            {
                (index - 1) as u64
            };

            uv_topology.last_mut().unwrap().push(uv_index);
            continue;
        }

        if !inner_tokens[1].is_empty()
        {
            let index = inner_tokens[1].parse::<i64>().unwrap();
            let uv_index = if index < 0
            {
                (uv_count as i64 + index) as u64
            }
            else
            {
                (index - 1) as u64
            };

            uv_topology.last_mut().unwrap().push(uv_index);
        }

        if !inner_tokens[2].is_empty()
        {
            let index = inner_tokens[2].parse::<i64>().unwrap();
            let normal_index = if index < 0
            {
                (normal_count as i64 + index) as u64
            }
            else
            {
                (index - 1) as u64
            };

            normal_topology.last_mut().unwrap().push(normal_index);
        }
    }

    // Prevent creating empty faces.
    if uv_topology.last().unwrap().is_empty()
    {
        uv_topology.pop();
    }
    if normal_topology.last().unwrap().is_empty()
    {
        normal_topology.pop();
    }
}

#[cfg(test)]
mod test
{
    use super::*;

    #[test]
    fn test_load_model()
    {
        let data = ObjData::from_disk_file("../../../Assets/cube_fs.obj");

        ObjData::export(&data, "wavefront_test.obj");
    }
}