path: root/engine/src/file_format/wavefront/obj.rs

//! OBJ file format parsing.
//!
//! File format documentation: <https://paulbourke.net/dataformats/obj>

use std::fs::read_to_string;
use std::path::PathBuf;

use crate::file_format::wavefront::common::{
    keyword,
    parse_statement_line,
    ParsingError,
    Statement,
    Triplet,
};
use crate::mesh::Mesh;
use crate::util::try_option;
use crate::vector::{Vec2, Vec3};
use crate::vertex::{Builder as VertexBuilder, Vertex};

/// Parses the content of a Wavefront `.obj`.
///
/// # Errors
/// Will return `Err` if the `.obj` content is formatted incorrectly.
pub fn parse(obj_content: &str) -> Result<Obj, Error>
{
    let lines = obj_content
        .lines()
        .enumerate()
        .map(|(line_index, line)| (line_index + 1, line));

    let statements = lines
        .map(|(line_no, line)| (line_no, parse_statement_line::<Keyword>(line, line_no)))
        .filter_map(|(line_no, result)| {
            let opt_statement = match result {
                Ok(opt_statement) => opt_statement,
                Err(err) => {
                    return Some(Err(err));
                }
            };

            Some(Ok((line_no, opt_statement?)))
        })
        .collect::<Result<Vec<_>, _>>()?;

    let vertex_positions = get_vertex_positions_from_statements(&statements)?;
    let texture_positions = get_texture_positions_from_statements(&statements)?;
    let vertex_normals = get_vertex_normals_from_statements(&statements)?;
    let material_specs = get_material_specs_from_statements(&statements)?;
    let faces = get_faces_from_statements(&statements, &material_specs)?;
    let mtl_libs = get_mtl_libs_from_statements(&statements)?;

    Ok(Obj {
        vertex_positions,
        vertex_normals,
        texture_positions,
        faces,
        mtl_libs,
    })
}

/// The data of a Wavefront object file.
#[derive(Debug)]
#[non_exhaustive]
pub struct Obj
{
    pub vertex_positions: Vec<Vec3<f32>>,
    pub vertex_normals: Vec<Vec3<f32>>,
    pub texture_positions: Vec<Vec2<f32>>,
    pub faces: Vec<Face>,
    pub mtl_libs: Vec<Vec<PathBuf>>,
}

impl Obj
{
    /// Creates a [`Mesh`] from this Wavefront object file data.
    ///
    /// # Errors
    /// Returns `Err` if:
    /// - A face's vertex position cannot be found
    /// - A face's texture position cannot be found
    /// - A face's vertex normal cannot be found
    /// - A face index does not fit in a [`u32`]
    pub fn to_mesh(&self) -> Result<Mesh, Error>
    {
        let vertices = self
            .faces
            .iter()
            .flat_map(|face| face.vertices.clone())
            .map(|face_vertex| face_vertex.to_vertex(self))
            .collect::<Result<Vec<_>, Error>>()?;

        Ok(Mesh::new(
            vertices,
            Some(
                self.faces
                    .iter()
                    .flat_map(|face| face.vertices.clone())
                    .enumerate()
                    .map(|(index, _)| {
                        u32::try_from(index).map_err(|_| Error::FaceIndexTooBig(index))
                    })
                    .collect::<Result<Vec<_>, _>>()?,
            ),
        ))
    }

    /// Reads and parses the material libraries of this `Obj`.
    ///
    /// # Errors
    /// Returns `Err` if:
    /// - Reading the contents of a material library fails
    /// - Parsing a material library fails
    pub fn read_and_parse_material_libs<ParsedMaterial, ParseError>(
        &self,
        parse_material_lib: impl Fn(&str) -> Result<Vec<ParsedMaterial>, ParseError>,
    ) -> Result<Vec<ParsedMaterial>, MaterialLibsError<ParseError>>
    {
        self.mtl_libs
            .iter()
            .flatten()
            .map(|mtl_lib| {
                Ok(parse_material_lib(&read_to_string(mtl_lib).map_err(
                    |err| MaterialLibsError::ReadingMaterialLibFailed {
                        source: err,
                        material_lib: mtl_lib.clone(),
                    },
                )?)?)
            })
            .flat_map(|res| match res {
                Ok(inner) => Either::A(inner.into_iter().map(Ok)),
                Err(err) => Either::B(vec![Err(err)].into_iter()),
            })
            .collect::<Result<Vec<_>, _>>()
    }
}

#[derive(Debug)]
#[non_exhaustive]
pub struct Face
{
    pub vertices: [FaceVertex; 3],
    pub material_name: Option<String>,
}

#[derive(Debug, Clone)]
pub struct FaceVertex
{
    pub position: u32,
    pub texture: Option<u32>,
    pub normal: Option<u32>,
}

impl FaceVertex
{
    /// Tries to convert this face vertex into a [`Vertex`].
    ///
    /// # Errors
    /// Returns `Err` if:
    /// - The face's vertex position cannot be found in the given [`Obj`]
    /// - The face's texture position cannot be found in the given [`Obj`]
    /// - The face's vertex normal cannot be found in the given [`Obj`]
    pub fn to_vertex(&self, obj: &Obj) -> Result<Vertex, Error>
    {
        let mut vertex_builder = VertexBuilder::default();

        let vertex_pos = *obj.vertex_positions.get(self.position as usize - 1).ok_or(
            Error::FaceVertexPositionNotFound { vertex_pos_index: self.position },
        )?;

        vertex_builder = vertex_builder.pos(vertex_pos);

        if let Some(face_vertex_texture) = self.texture {
            let texture_pos = obj
                .texture_positions
                .get(face_vertex_texture as usize - 1)
                .ok_or(Error::FaceTexturePositionNotFound {
                    texture_pos_index: face_vertex_texture,
                })?;

            vertex_builder = vertex_builder.texture_coords(*texture_pos);
        }

        if let Some(face_vertex_normal) = self.normal {
            let vertex_normal = *obj
                .vertex_normals
                .get(face_vertex_normal as usize - 1)
                .ok_or(Error::FaceVertexNormalNotFound {
                    vertex_normal_index: face_vertex_normal,
                })?;

            vertex_builder = vertex_builder.normal(vertex_normal);
        }

        Ok(vertex_builder.build())
    }
}

impl From<Triplet> for FaceVertex
{
    fn from(triplet: Triplet) -> Self
    {
        Self {
            position: triplet.0,
            texture: triplet.1,
            normal: triplet.2,
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum Error
{
    #[error(transparent)]
    ParsingError(#[from] ParsingError),

    #[error(
        "Face vertex position with index {vertex_pos_index} (1-based) was not found"
    )]
    FaceVertexPositionNotFound
    {
        vertex_pos_index: u32
    },

    #[error(
        "Face texture position with index {texture_pos_index} (1-based) was not found"
    )]
    FaceTexturePositionNotFound
    {
        texture_pos_index: u32
    },

    #[error(
        "Face vertex normal with index {vertex_normal_index} (1-based) was not found"
    )]
    FaceVertexNormalNotFound
    {
        vertex_normal_index: u32
    },

    #[error("Face index {0} is too big to fit into a 32-bit integer")]
    FaceIndexTooBig(usize),

    #[error(
        "Unsupported number of arguments ({arg_count}) to {keyword} at line {line_no}"
    )]
    UnsupportedArgumentCount
    {
        keyword: String,
        arg_count: usize,
        line_no: usize,
    },

    #[error("Invalid number of arguments ({arg_count}) to {keyword} at line {line_no}")]
    InvalidArgumentCount
    {
        keyword: String,
        arg_count: usize,
        line_no: usize,
    },
}

#[derive(Debug, thiserror::Error)]
pub enum MaterialLibsError<MaterialLibParseError>
{
    #[error("Parsing material library failed")]
    ParsingFailed(#[from] MaterialLibParseError),

    #[error("Failed to read material library {}", material_lib.display())]
    ReadingMaterialLibFailed
    {
        #[source]
        source: std::io::Error,
        material_lib: PathBuf,
    },
}

keyword! {
    #[derive(Debug, PartialEq, Eq, Clone, Copy)]
    enum Keyword {
        #[keyword(rename = "v")]
        V,

        #[keyword(rename = "vn")]
        Vn,

        #[keyword(rename = "vt")]
        Vt,

        #[keyword(rename = "o")]
        O,

        #[keyword(rename = "s")]
        S,

        #[keyword(rename = "f")]
        F,

        #[keyword(rename = "mtllib")]
        Mtllib,

        #[keyword(rename = "usemtl")]
        Usemtl,
    }
}

#[derive(Debug)]
struct MaterialSpecifier
{
    material_name: String,
    line_no: usize,
}

fn find_material_specifier_for_line_no(
    material_specifiers: &[MaterialSpecifier],
    line_no: usize,
) -> Option<&MaterialSpecifier>
{
    material_specifiers
        .iter()
        .rev()
        .find(|material_specifier| material_specifier.line_no < line_no)
}

enum Either<ValA, ValB>
{
    A(ValA),
    B(ValB),
}

impl<ValA, ValB, Item> Iterator for Either<ValA, ValB>
where
    ValA: Iterator<Item = Item>,
    ValB: Iterator<Item = Item>,
{
    type Item = Item;

    fn next(&mut self) -> Option<Self::Item>
    {
        match self {
            Self::A(iter_a) => iter_a.next(),
            Self::B(iter_b) => iter_b.next(),
        }
    }
}

fn get_vertex_positions_from_statements(
    statements: &[(usize, Statement<Keyword>)],
) -> Result<Vec<Vec3<f32>>, Error>
{
    statements
        .iter()
        .filter_map(|(line_no, statement)| {
            if statement.keyword != Keyword::V {
                return None;
            }

            if statement.arguments.len() == 4 {
                return Some(Err(Error::UnsupportedArgumentCount {
                    keyword: statement.keyword.to_string(),
                    arg_count: statement.arguments.len(),
                    line_no: *line_no,
                }));
            }

            if statement.arguments.len() > 4 {
                return Some(Err(Error::InvalidArgumentCount {
                    keyword: statement.keyword.to_string(),
                    arg_count: statement.arguments.len(),
                    line_no: *line_no,
                }));
            }

            let x = try_option!(statement.get_float_arg(0, *line_no));
            let y = try_option!(statement.get_float_arg(1, *line_no));
            let z = try_option!(statement.get_float_arg(2, *line_no));

            Some(Ok(Vec3 { x, y, z }))
        })
        .collect::<Result<Vec<_>, Error>>()
}

fn get_texture_positions_from_statements(
    statements: &[(usize, Statement<Keyword>)],
) -> Result<Vec<Vec2<f32>>, Error>
{
    statements
        .iter()
        .filter_map(|(line_no, statement)| {
            if statement.keyword != Keyword::Vt {
                return None;
            }

            if statement.arguments.len() == 3 {
                return Some(Err(Error::UnsupportedArgumentCount {
                    keyword: statement.keyword.to_string(),
                    arg_count: statement.arguments.len(),
                    line_no: *line_no,
                }));
            }

            if statement.arguments.len() > 3 {
                return Some(Err(Error::InvalidArgumentCount {
                    keyword: statement.keyword.to_string(),
                    arg_count: statement.arguments.len(),
                    line_no: *line_no,
                }));
            }

            let u = try_option!(statement.get_float_arg(0, *line_no));
            let v = try_option!(statement.get_float_arg(1, *line_no));

            Some(Ok(Vec2 { x: u, y: v }))
        })
        .collect::<Result<Vec<_>, Error>>()
}

fn get_vertex_normals_from_statements(
    statements: &[(usize, Statement<Keyword>)],
) -> Result<Vec<Vec3<f32>>, Error>
{
    statements
        .iter()
        .filter_map(|(line_no, statement)| {
            if statement.keyword != Keyword::Vn {
                return None;
            }

            if statement.arguments.len() > 3 {
                return Some(Err(Error::InvalidArgumentCount {
                    keyword: statement.keyword.to_string(),
                    arg_count: statement.arguments.len(),
                    line_no: *line_no,
                }));
            }

            let i = try_option!(statement.get_float_arg(0, *line_no));
            let j = try_option!(statement.get_float_arg(1, *line_no));
            let k = try_option!(statement.get_float_arg(2, *line_no));

            Some(Ok(Vec3 { x: i, y: j, z: k }))
        })
        .collect::<Result<Vec<_>, Error>>()
}

fn get_material_specs_from_statements(
    statements: &[(usize, Statement<Keyword>)],
) -> Result<Vec<MaterialSpecifier>, Error>
{
    statements
        .iter()
        .filter_map(|(line_no, statement)| {
            if statement.keyword != Keyword::Usemtl {
                return None;
            }

            if statement.arguments.len() > 1 {
                return Some(Err(Error::InvalidArgumentCount {
                    keyword: statement.keyword.to_string(),
                    arg_count: statement.arguments.len(),
                    line_no: *line_no,
                }));
            }

            let material_name = try_option!(statement.get_text_arg(0, *line_no));

            Some(Ok(MaterialSpecifier {
                material_name: material_name.to_string(),
                line_no: *line_no,
            }))
        })
        .collect::<Result<Vec<_>, Error>>()
}

fn get_faces_from_statements(
    statements: &[(usize, Statement<Keyword>)],
    material_specifiers: &[MaterialSpecifier],
) -> Result<Vec<Face>, Error>
{
    statements
        .iter()
        .filter_map(|(line_no, statement)| {
            if statement.keyword != Keyword::F {
                return None;
            }

            if statement.arguments.len() > 3 {
                return Some(Err(Error::UnsupportedArgumentCount {
                    keyword: statement.keyword.to_string(),
                    arg_count: statement.arguments.len(),
                    line_no: *line_no,
                }));
            }

            let vertex_a = try_option!(statement.get_triplet_arg(0, *line_no)).into();
            let vertex_b = try_option!(statement.get_triplet_arg(1, *line_no)).into();
            let vertex_c = try_option!(statement.get_triplet_arg(2, *line_no)).into();

            let material_name =
                find_material_specifier_for_line_no(material_specifiers, *line_no)
                    .map(|material_specifier| material_specifier.material_name.clone());

            Some(Ok(Face {
                vertices: [vertex_a, vertex_b, vertex_c],
                material_name,
            }))
        })
        .collect::<Result<Vec<Face>, Error>>()
}

fn get_mtl_libs_from_statements(
    statements: &[(usize, Statement<Keyword>)],
) -> Result<Vec<Vec<PathBuf>>, Error>
{
    statements
        .iter()
        .filter_map(|(line_no, statement)| {
            if statement.keyword != Keyword::Mtllib {
                return None;
            }

            let mtl_lib_paths = try_option!(statement
                .arguments
                .iter()
                .enumerate()
                .map(|(index, value)| Ok(PathBuf::from(value.to_text(index, *line_no)?)))
                .collect::<Result<Vec<_>, ParsingError>>());

            Some(Ok(mtl_lib_paths))
        })
        .collect::<Result<Vec<_>, Error>>()
}

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

    #[test]
    fn parse_containing_usemtl_works()
    {
        let obj = parse(concat!(
            "usemtl dark-green\n",
            "f 6/8/1 7/5/3 1/3/2\n",
            "f 1/2/4 4/1/8 3/4/2\n",
            "usemtl light-pink\n",
            "f 9/1/7 5/1/2 3/2/7"
        ))
        .expect("Expected Ok");

        assert_eq!(obj.faces.len(), 3);

        assert_eq!(
            obj.faces[0].material_name.as_ref().expect("Expected Some"),
            "dark-green"
        );

        assert_eq!(
            obj.faces[1].material_name.as_ref().expect("Expected Some"),
            "dark-green"
        );

        assert_eq!(
            obj.faces[2].material_name.as_ref().expect("Expected Some"),
            "light-pink"
        );
    }
}