draco_io/

obj_writer.rs

//! OBJ format writer for meshes and point clouds.
//!
//! Supports writing:
//! - Vertex positions
//! - Triangle faces (for meshes)
//! - Vertex normals (if present)
//! - Vertex texture coordinates (if present)
//!
//! # Example
//!
//! ```ignore
//! use draco_io::ObjWriter;
//! use draco_core::mesh::Mesh;
//!
//! let mesh: Mesh = /* ... */;
//! let mut writer = ObjWriter::new();
//! writer.add_mesh(&mesh, Some("MyMesh"));
//! writer.write("output.obj")?;
//!
//! // Or write point cloud
//! let mut writer = ObjWriter::new();
//! writer.add_points(&[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]);
//! writer.write("points.obj")?;
//! ```

use std::fs::File;
use std::io::{self, BufWriter, Write};
use std::path::Path;

use draco_core::draco_types::DataType;
use draco_core::geometry_attribute::GeometryAttributeType;
use draco_core::geometry_indices::FaceIndex;
use draco_core::mesh::Mesh;

use crate::traits::{PointCloudWriter, WriteToBytes, Writer};

/// OBJ format writer.
///
/// This struct provides a builder-style API for writing OBJ files.
/// Meshes or points are added, then written with `write()`.
///
/// # Example
///
/// ```ignore
/// use draco_io::ObjWriter;
///
/// let mut writer = ObjWriter::new();
/// writer.add_mesh(&mesh, Some("Cube"));
/// writer.write("cube.obj")?;
/// ```
#[derive(Debug, Clone, Default)]
pub struct ObjWriter {
    /// Collected vertex positions
    positions: Vec<[f32; 3]>,
    /// Collected vertex normals
    normals: Vec<[f32; 3]>,
    /// Collected vertex texture coordinates
    texcoords: Vec<[f32; 2]>,
    /// Collected faces (1-based indices)
    faces: Vec<ObjFace>,
    /// Object groups with (name, start_face_index)
    groups: Vec<(String, usize)>,
}

#[derive(Debug, Clone)]
struct ObjFace {
    positions: [u32; 3],
    texcoords: Option<[u32; 3]>,
    normals: Option<[u32; 3]>,
}

impl ObjWriter {
    /// Create a new OBJ writer.
    pub fn new() -> Self {
        Self::default()
    }

    /// Add raw point positions (for point cloud output).
    pub fn add_points(&mut self, points: &[[f32; 3]]) {
        self.positions.extend_from_slice(points);
    }

    /// Add a single point.
    pub fn add_point(&mut self, point: [f32; 3]) {
        self.positions.push(point);
    }

    /// Get the number of vertices added.
    pub fn vertex_count(&self) -> usize {
        self.positions.len()
    }

    /// Get the number of faces added.
    pub fn face_count(&self) -> usize {
        self.faces.len()
    }

    /// Write the OBJ file to the given path.
    pub fn write<P: AsRef<Path>>(&self, path: P) -> io::Result<()> {
        let file = File::create(path)?;
        let mut writer = BufWriter::new(file);
        self.write_to(&mut writer)
    }

    /// Write the OBJ data to a writer.
    pub fn write_to<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        // Header comment
        writeln!(writer, "# OBJ file generated by draco-io")?;
        writeln!(writer, "# Vertices: {}", self.positions.len())?;
        writeln!(writer, "# Faces: {}", self.faces.len())?;
        writeln!(writer)?;

        // Write all positions
        for [x, y, z] in &self.positions {
            writeln!(writer, "v {:.6} {:.6} {:.6}", x, y, z)?;
        }

        // Write texture coordinates if present
        if !self.texcoords.is_empty() {
            writeln!(writer)?;
            for [u, v] in &self.texcoords {
                writeln!(writer, "vt {:.6} {:.6}", u, v)?;
            }
        }

        // Write normals if present
        if !self.normals.is_empty() {
            writeln!(writer)?;
            for [x, y, z] in &self.normals {
                writeln!(writer, "vn {:.6} {:.6} {:.6}", x, y, z)?;
            }
        }

        // Write faces with groups
        if !self.faces.is_empty() {
            writeln!(writer)?;

            let mut group_iter = self.groups.iter().peekable();
            let has_texcoords = !self.texcoords.is_empty();
            let has_normals = !self.normals.is_empty();

            for (i, face) in self.faces.iter().enumerate() {
                // Check if we need to start a new group
                if let Some((name, start_idx)) = group_iter.peek() {
                    if i == *start_idx {
                        writeln!(writer, "o {}", name)?;
                        group_iter.next();
                    }
                }

                match (
                    has_texcoords.then_some(face.texcoords).flatten(),
                    has_normals.then_some(face.normals).flatten(),
                ) {
                    (Some(texcoords), Some(normals)) => {
                        writeln!(
                            writer,
                            "f {}/{}/{} {}/{}/{} {}/{}/{}",
                            face.positions[0],
                            texcoords[0],
                            normals[0],
                            face.positions[1],
                            texcoords[1],
                            normals[1],
                            face.positions[2],
                            texcoords[2],
                            normals[2]
                        )?;
                    }
                    (Some(texcoords), None) => {
                        writeln!(
                            writer,
                            "f {}/{} {}/{} {}/{}",
                            face.positions[0],
                            texcoords[0],
                            face.positions[1],
                            texcoords[1],
                            face.positions[2],
                            texcoords[2]
                        )?;
                    }
                    (None, Some(normals)) => {
                        writeln!(
                            writer,
                            "f {}//{} {}//{} {}//{}",
                            face.positions[0],
                            normals[0],
                            face.positions[1],
                            normals[1],
                            face.positions[2],
                            normals[2]
                        )?;
                    }
                    (None, None) => {
                        writeln!(
                            writer,
                            "f {} {} {}",
                            face.positions[0], face.positions[1], face.positions[2]
                        )?;
                    }
                }
            }
        }

        Ok(())
    }

    /// Write the OBJ data into a byte vector.
    pub fn write_to_vec(&self) -> io::Result<Vec<u8>> {
        let mut out = Vec::new();
        self.write_to(&mut out)?;
        Ok(out)
    }
}

/// Read a float3 from an attribute at a given point index.
fn read_float3(mesh: &Mesh, att_id: i32, point_idx: usize) -> [f32; 3] {
    let att = mesh.attribute(att_id);
    let byte_stride = att.byte_stride() as usize;
    let buffer = att.buffer();
    let mut bytes = [0u8; 12];
    buffer.read(point_idx * byte_stride, &mut bytes);
    [
        f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]),
        f32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]),
        f32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]),
    ]
}

/// Read a float2 from an attribute at a given point index.
fn read_float2(mesh: &Mesh, att_id: i32, point_idx: usize) -> [f32; 2] {
    let att = mesh.attribute(att_id);
    let byte_stride = att.byte_stride() as usize;
    let buffer = att.buffer();
    let mut bytes = [0u8; 8];
    buffer.read(point_idx * byte_stride, &mut bytes);
    [
        f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]),
        f32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]),
    ]
}

fn require_f32_components(mesh: &Mesh, att_id: i32, components: u8, label: &str) -> io::Result<()> {
    let att = mesh.attribute(att_id);
    if att.data_type() != DataType::Float32 || att.num_components() != components {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("OBJ writer requires {label} attributes to be Float32x{components}"),
        ));
    }
    Ok(())
}

// ============================================================================
// Trait Implementations
// ============================================================================

impl Writer for ObjWriter {
    fn new() -> Self {
        Self::default()
    }

    fn add_mesh(&mut self, mesh: &Mesh, name: Option<&str>) -> io::Result<()> {
        let vertex_offset = self.positions.len() as u32;
        let normal_offset = self.normals.len() as u32;
        let texcoord_offset = self.texcoords.len() as u32;
        let face_start = self.faces.len();

        // Add group if name provided
        if let Some(n) = name {
            self.groups.push((n.to_string(), face_start));
        }

        // Extract positions
        let pos_att_id = mesh.named_attribute_id(GeometryAttributeType::Position);
        if pos_att_id >= 0 {
            require_f32_components(mesh, pos_att_id, 3, "position")?;
            for i in 0..mesh.num_points() {
                self.positions.push(read_float3(mesh, pos_att_id, i));
            }
        }

        // Extract normals if present
        let normal_att_id = mesh.named_attribute_id(GeometryAttributeType::Normal);
        let has_normals = normal_att_id >= 0;
        if has_normals {
            require_f32_components(mesh, normal_att_id, 3, "normal")?;
            for i in 0..mesh.num_points() {
                self.normals.push(read_float3(mesh, normal_att_id, i));
            }
        }

        // Extract texture coordinates if present
        let texcoord_att_id = mesh.named_attribute_id(GeometryAttributeType::TexCoord);
        let has_texcoords = texcoord_att_id >= 0;
        if has_texcoords {
            require_f32_components(mesh, texcoord_att_id, 2, "texcoord")?;
            for i in 0..mesh.num_points() {
                self.texcoords.push(read_float2(mesh, texcoord_att_id, i));
            }
        }

        // Extract faces (convert to 1-based indices with offsets)
        for i in 0..mesh.num_faces() as u32 {
            let face = mesh.face(FaceIndex(i));
            let positions = [
                face[0].0 + vertex_offset + 1,
                face[1].0 + vertex_offset + 1,
                face[2].0 + vertex_offset + 1,
            ];
            let normals = has_normals.then(|| {
                [
                    face[0].0 + normal_offset + 1,
                    face[1].0 + normal_offset + 1,
                    face[2].0 + normal_offset + 1,
                ]
            });
            let texcoords = has_texcoords.then(|| {
                [
                    face[0].0 + texcoord_offset + 1,
                    face[1].0 + texcoord_offset + 1,
                    face[2].0 + texcoord_offset + 1,
                ]
            });
            self.faces.push(ObjFace {
                positions,
                texcoords,
                normals,
            });
        }
        Ok(())
    }

    fn write<P: AsRef<Path>>(&self, path: P) -> io::Result<()> {
        self.write(path)
    }

    fn vertex_count(&self) -> usize {
        self.vertex_count()
    }

    fn face_count(&self) -> usize {
        self.face_count()
    }
}

impl PointCloudWriter for ObjWriter {
    fn add_points(&mut self, points: &[[f32; 3]]) {
        self.positions.extend_from_slice(points);
    }

    fn add_point(&mut self, point: [f32; 3]) {
        self.positions.push(point);
    }
}

impl WriteToBytes for ObjWriter {
    fn write_to_vec(&self) -> io::Result<Vec<u8>> {
        ObjWriter::write_to_vec(self)
    }

    fn write_to<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        ObjWriter::write_to(self, writer)
    }
}

// ============================================================================
// Convenience Functions (for backward compatibility)
// ============================================================================

/// Write a mesh to an OBJ file with positions and faces.
///
/// This is a convenience function. For more control, use `ObjWriter` directly.
pub fn write_obj_mesh<P: AsRef<Path>>(path: P, mesh: &Mesh) -> io::Result<()> {
    let mut writer = ObjWriter::new();
    Writer::add_mesh(&mut writer, mesh, None)?;
    writer.write(path)
}

/// Write point positions to an OBJ file (point cloud, no faces).
///
/// This is a convenience function. For more control, use `ObjWriter` directly.
pub fn write_obj_positions<P: AsRef<Path>>(path: P, points: &[[f32; 3]]) -> io::Result<()> {
    let mut writer = ObjWriter::new();
    writer.add_points(points);
    writer.write(path)
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use draco_core::draco_types::DataType;
    use draco_core::geometry_attribute::PointAttribute;
    use draco_core::geometry_indices::PointIndex;
    use std::fs;
    use std::io::{BufRead, BufReader};
    use tempfile::NamedTempFile;

    fn create_triangle_mesh() -> Mesh {
        let mut mesh = Mesh::new();
        let mut pos_att = PointAttribute::new();

        pos_att.init(
            GeometryAttributeType::Position,
            3,
            DataType::Float32,
            false,
            3,
        );
        let buffer = pos_att.buffer_mut();
        let positions: [[f32; 3]; 3] = [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]];
        for (i, pos) in positions.iter().enumerate() {
            let bytes: Vec<u8> = pos.iter().flat_map(|v| v.to_le_bytes()).collect();
            buffer.write(i * 12, &bytes);
        }
        mesh.add_attribute(pos_att);

        mesh.set_num_faces(1);
        mesh.set_face(FaceIndex(0), [PointIndex(0), PointIndex(1), PointIndex(2)]);

        mesh
    }

    fn add_f32_attribute(
        mesh: &mut Mesh,
        attribute_type: GeometryAttributeType,
        components: u8,
        values: &[f32],
    ) {
        let mut att = PointAttribute::new();
        att.init(
            attribute_type,
            components,
            DataType::Float32,
            false,
            values.len() / components as usize,
        );
        let bytes: Vec<u8> = values
            .iter()
            .flat_map(|component| component.to_le_bytes())
            .collect();
        att.buffer_mut().write(0, &bytes);
        mesh.add_attribute(att);
    }

    #[test]
    fn test_obj_writer_new() {
        let writer = ObjWriter::new();
        assert_eq!(writer.vertex_count(), 0);
        assert_eq!(writer.face_count(), 0);
    }

    #[test]
    fn test_obj_writer_add_mesh() {
        let mesh = create_triangle_mesh();
        let mut writer = ObjWriter::new();
        Writer::add_mesh(&mut writer, &mesh, Some("Triangle")).unwrap();
        assert_eq!(writer.vertex_count(), 3);
        assert_eq!(writer.face_count(), 1);
    }

    #[test]
    fn test_obj_writer_add_points() {
        let mut writer = ObjWriter::new();
        writer.add_points(&[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
        assert_eq!(writer.vertex_count(), 2);
        assert_eq!(writer.face_count(), 0);
    }

    #[test]
    fn test_write_obj_positions() {
        let points = vec![[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]];

        let file = NamedTempFile::new().unwrap();
        write_obj_positions(file.path(), &points).unwrap();

        let content = fs::read_to_string(file.path()).unwrap();
        assert!(content.contains("v 0.000000 0.000000 0.000000"));
        assert!(content.contains("v 1.000000 0.000000 0.000000"));
        assert!(content.contains("v 0.000000 1.000000 0.000000"));
    }

    #[test]
    fn test_write_obj_mesh() {
        let mesh = create_triangle_mesh();
        let file = NamedTempFile::new().unwrap();
        write_obj_mesh(file.path(), &mesh).unwrap();

        let reader = BufReader::new(fs::File::open(file.path()).unwrap());
        let lines: Vec<String> = reader.lines().map_while(Result::ok).collect();

        // Check vertices and face
        assert!(lines.iter().any(|l| l.starts_with("v ")));
        assert!(lines.iter().any(|l| l == "f 1 2 3"));
    }

    #[test]
    fn test_multiple_meshes() {
        let mesh1 = create_triangle_mesh();
        let mesh2 = create_triangle_mesh();

        let mut writer = ObjWriter::new();
        Writer::add_mesh(&mut writer, &mesh1, Some("Mesh1")).unwrap();
        Writer::add_mesh(&mut writer, &mesh2, Some("Mesh2")).unwrap();

        assert_eq!(writer.vertex_count(), 6);
        assert_eq!(writer.face_count(), 2);

        let file = NamedTempFile::new().unwrap();
        writer.write(file.path()).unwrap();

        let content = fs::read_to_string(file.path()).unwrap();
        assert!(content.contains("o Mesh1"));
        assert!(content.contains("o Mesh2"));
        // Second mesh should have offset indices
        assert!(content.contains("f 4 5 6"));
    }

    #[test]
    fn test_write_obj_mesh_with_normals_and_texcoords() {
        let mut mesh = create_triangle_mesh();
        add_f32_attribute(
            &mut mesh,
            GeometryAttributeType::Normal,
            3,
            &[0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0],
        );
        add_f32_attribute(
            &mut mesh,
            GeometryAttributeType::TexCoord,
            2,
            &[0.0, 0.0, 1.0, 0.0, 0.0, 1.0],
        );

        let mut writer = ObjWriter::new();
        Writer::add_mesh(&mut writer, &mesh, Some("Triangle")).unwrap();

        let content = String::from_utf8(writer.write_to_vec().unwrap()).unwrap();
        assert!(content.contains("vt 1.000000 0.000000"));
        assert!(content.contains("vn 0.000000 0.000000 1.000000"));
        assert!(content.contains("f 1/1/1 2/2/2 3/3/3"));
    }
}