pbrt/src/shapes/mesh.rs

use anyhow::{bail, Context, Result as AnyResult};
use ply_rs::parser::Parser;
use ply_rs::ply::{DefaultElement, Property};
use shared::core::geometry::{Normal3f, Point2f, Point3f, VectorLike};
use shared::shapes::mesh::TriangleMesh;
use shared::utils::sampling::PiecewiseConstant2D;
use shared::Transform;
use std::fs::File;
use std::path::Path;

#[derive(Debug, Clone, Default)]
pub struct TriQuadMesh {
    pub p: Vec<Point3f>,
    pub n: Vec<Normal3f>,
    pub uv: Vec<Point2f>,
    pub face_indices: Vec<i32>,
    pub tri_indices: Vec<i32>,
    pub quad_indices: Vec<i32>,
}

// PLY Helper Functions
fn get_float(elem: &DefaultElement, key: &str) -> AnyResult<f32> {
    match elem.get(key) {
        Some(Property::Float(v)) => Ok(*v),
        Some(Property::Double(v)) => Ok(*v as f32),
        Some(_) => bail!("Property {} is not a float", key),
        None => bail!("Property {} not found", key),
    }
}

fn get_int(elem: &DefaultElement, key: &str) -> AnyResult<i32> {
    match elem.get(key) {
        Some(Property::Int(v)) => Ok(*v),
        Some(Property::UInt(v)) => Ok(*v as i32),
        Some(Property::Short(v)) => Ok(*v as i32),
        Some(Property::UShort(v)) => Ok(*v as i32),
        Some(Property::Char(v)) => Ok(*v as i32),
        Some(Property::UChar(v)) => Ok(*v as i32),
        _ => bail!("Property {} not found or not integer", key),
    }
}

fn get_float_any(elem: &DefaultElement, keys: &[&str]) -> Option<f32> {
    for k in keys {
        if let Ok(val) = get_float(elem, k) {
            return Some(val);
        }
    }
    None
}

fn get_list_uint(elem: &DefaultElement, key: &str) -> AnyResult<Vec<u32>> {
    match elem.get(key) {
        Some(Property::ListInt(vec)) => Ok(vec.iter().map(|&x| x as u32).collect()),
        Some(Property::ListUInt(vec)) => Ok(vec.clone()),
        Some(Property::ListUChar(vec)) => Ok(vec.iter().map(|&x| x as u32).collect()),
        Some(Property::ListChar(vec)) => Ok(vec.iter().map(|&x| x as u32).collect()),
        _ => bail!("Property {} is not a list", key),
    }
}

impl TriQuadMesh {
    pub fn read_ply<P: AsRef<Path>>(filename: P) -> AnyResult<Self> {
        let path = filename.as_ref();
        let filename_display = path.display().to_string();

        let f = File::open(path)
            .with_context(|| format!("Couldn't open PLY file \"{}\"", filename_display))?;

        let p = Parser::<DefaultElement>::new();
        let ply = if path.extension().and_then(|e| e.to_str()) == Some("gz") {
            let decoder = flate2::read::GzDecoder::new(f);
            let mut buf = std::io::BufReader::new(decoder);
            p.read_ply(&mut buf)
        } else {
            let mut buf = std::io::BufReader::new(f);
            p.read_ply(&mut buf)
        }
        .with_context(|| format!("Unable to read/parse PLY file \"{}\"", filename_display))?;

        let mut mesh = TriQuadMesh::default();

        if let Some(vertices) = ply.payload.get("vertex") {
            if vertices.is_empty() {
                bail!("{}: PLY file has no vertices", filename_display);
            }

            let first = &vertices[0];
            let has_uv = (first.contains_key("u") && first.contains_key("v"))
                || (first.contains_key("s") && first.contains_key("t"))
                || (first.contains_key("texture_u") && first.contains_key("texture_v"))
                || (first.contains_key("texture_s") && first.contains_key("texture_t"));
            let has_normal =
                first.contains_key("nx") && first.contains_key("ny") && first.contains_key("nz");

            mesh.p.reserve(vertices.len());
            if has_normal {
                mesh.n.reserve(vertices.len());
            }
            if has_uv {
                mesh.uv.reserve(vertices.len());
            }

            for v_elem in vertices {
                let x = get_float(v_elem, "x")?;
                let y = get_float(v_elem, "y")?;
                let z = get_float(v_elem, "z")?;
                mesh.p.push(Point3f::new(x, y, z));

                if has_normal {
                    let nx = get_float(v_elem, "nx").unwrap_or(0.0);
                    let ny = get_float(v_elem, "ny").unwrap_or(0.0);
                    let nz = get_float(v_elem, "nz").unwrap_or(0.0);
                    mesh.n.push(Normal3f::new(nx, ny, nz));
                }

                if has_uv {
                    let u =
                        get_float_any(v_elem, &["u", "s", "texture_u", "texture_s"]).unwrap_or(0.0);
                    let v =
                        get_float_any(v_elem, &["v", "t", "texture_v", "texture_t"]).unwrap_or(0.0);
                    mesh.uv.push(Point2f::new(u, v));
                }
            }
        } else {
            bail!("{}: PLY file has no vertex elements", filename_display);
        }

        if let Some(faces) = ply.payload.get("face") {
            mesh.tri_indices.reserve(faces.len() * 3);
            mesh.quad_indices.reserve(faces.len() * 4);

            for f_elem in faces {
                if let Ok(fi) = get_int(f_elem, "face_indices") {
                    mesh.face_indices.push(fi);
                }

                if let Ok(indices) = get_list_uint(f_elem, "vertex_indices") {
                    match indices.len() {
                        3 => mesh.tri_indices.extend(indices.iter().map(|&i| i as i32)),
                        4 => mesh.quad_indices.extend(indices.iter().map(|&i| i as i32)),
                        n => {
                            log::warn!(
                                "{}: Skipping face with {} vertices (only 3 or 4 supported)",
                                filename_display,
                                n
                            );
                        }
                    }
                } else {
                    bail!("{}: vertex_indices not found in face", filename_display);
                }
            }
        } else {
            bail!("{}: PLY file has no face elements", filename_display);
        }

        let vertex_count = mesh.p.len() as i32;
        for &idx in &mesh.tri_indices {
            if idx >= vertex_count {
                bail!(
                    "{}: Vertex index {} out of bounds [0..{})",
                    filename_display,
                    idx,
                    vertex_count
                );
            }
        }
        for &idx in &mesh.quad_indices {
            if idx >= vertex_count {
                bail!(
                    "{}: Vertex index {} out of bounds [0..{})",
                    filename_display,
                    idx,
                    vertex_count
                );
            }
        }

        Ok(mesh)
    }

    pub fn convert_to_only_triangles(&mut self) {
        if self.quad_indices.is_empty() {
            return;
        }

        self.tri_indices.reserve(self.quad_indices.len() / 4 * 6);

        for quad in self.quad_indices.chunks_exact(4) {
            let (v0, v1, v2, v3) = (quad[0], quad[1], quad[2], quad[3]);
            self.tri_indices.push(v0);
            self.tri_indices.push(v1);
            self.tri_indices.push(v2);
            self.tri_indices.push(v0);
            self.tri_indices.push(v2);
            self.tri_indices.push(v3);
        }

        self.quad_indices.clear();
    }

    pub fn compute_normals(&mut self) {
        self.n = vec![Normal3f::new(0.0, 0.0, 0.0); self.p.len()];

        for tri in self.tri_indices.chunks_exact(3) {
            let (i0, i1, i2) = (tri[0] as usize, tri[1] as usize, tri[2] as usize);

            let p0 = self.p[i0];
            let p1 = self.p[i1];
            let p2 = self.p[i2];

            let v10 = p1 - p0;
            let v20 = p2 - p0;
            let face_normal = v10.cross(v20);

            self.n[i0] = self.n[i0] + Normal3f::from(face_normal);
            self.n[i1] = self.n[i1] + Normal3f::from(face_normal);
            self.n[i2] = self.n[i2] + Normal3f::from(face_normal);
        }

        for quad in self.quad_indices.chunks_exact(4) {
            let indices: Vec<usize> = quad.iter().map(|&i| i as usize).collect();

            let p0 = self.p[indices[0]];
            let p1 = self.p[indices[1]];
            let p2 = self.p[indices[2]];

            let v10 = p1 - p0;
            let v20 = p2 - p0;
            let face_normal = v10.cross(v20);

            for &idx in &indices {
                self.n[idx] = self.n[idx] + Normal3f::from(face_normal);
            }
        }

        for normal in &mut self.n {
            let len_sq = normal.norm_squared();
            if len_sq > 0.0 {
                *normal = *normal / len_sq.sqrt();
            }
        }
    }

    pub fn into_triangle_mesh(
        mut self,
        render_from_object: &Transform,
        reverse_orientation: bool,
    ) -> TriangleMesh {
        self.convert_to_only_triangles();

        if self.n.is_empty() {
            self.compute_normals();
        }

        TriangleMesh::new(
            render_from_object,
            reverse_orientation,
            &self.tri_indices,
            &self.p,
            &self.n,
            &Vec::new(),
            &self.uv,
            &self.face_indices,
        )
    }
}

pub trait ReadTriangleMesh {
    fn from_ply<P: AsRef<Path>>(
        filename: P,
        render_from_object: &Transform,
        reverse_orientation: bool,
    ) -> AnyResult<Self>
    where
        Self: Sized;
}

impl ReadTriangleMesh for TriangleMesh {
    fn from_ply<P: AsRef<Path>>(
        filename: P,
        render_from_object: &Transform,
        reverse_orientation: bool,
    ) -> AnyResult<Self> {
        let mesh = TriQuadMesh::read_ply(filename)?;
        Ok(mesh.into_triangle_mesh(render_from_object, reverse_orientation))
    }
}