openctm 0.1.0

//! Pure-Rust decoder for the [OpenCTM](https://openctm.sourceforge.net/) mesh
//! format (MG1 method), the geometry encoding used by ContextCapture 3mx.
//!
//! ```no_run
//! let mesh = openctm::decode(&std::fs::read("mesh.ctm").unwrap()).unwrap();
//! ```

use std::io::Cursor;

const FORMAT_VERSION: u32 = 5;
const HAS_NORMALS_BIT: u32 = 0x1;
const MAX_ELEMS: usize = 64 << 20;

pub struct CtmMesh {
    pub positions: Vec<[f32; 3]>,
    pub normals: Option<Vec<[f32; 3]>>,
    pub uvs: Option<Vec<[f32; 2]>>,
    pub indices: Vec<u32>,
}

struct Reader<'a> {
    d: &'a [u8],
    pos: usize,
}

impl<'a> Reader<'a> {
    fn bytes(&mut self, n: usize) -> Result<&'a [u8], String> {
        let end = self.pos.checked_add(n).ok_or("overflow")?;
        if end > self.d.len() {
            return Err("unexpected end of CTM data".into());
        }
        let s = &self.d[self.pos..end];
        self.pos = end;
        Ok(s)
    }

    fn u32(&mut self) -> Result<u32, String> {
        let b = self.bytes(4)?;
        Ok(u32::from_le_bytes([b[0], b[1], b[2], b[3]]))
    }

    fn magic(&mut self, expect: &[u8; 4]) -> Result<(), String> {
        let b = self.bytes(4)?;
        if b != expect {
            return Err(format!(
                "expected {:?}, got {:?}",
                std::str::from_utf8(expect).unwrap_or("?"),
                String::from_utf8_lossy(b)
            ));
        }
        Ok(())
    }

    fn skip_string(&mut self) -> Result<(), String> {
        let len = self.u32()? as usize;
        self.bytes(len)?;
        Ok(())
    }

    fn packed_words(&mut self, count: usize, size: usize) -> Result<Vec<u32>, String> {
        let elems = count
            .checked_mul(size)
            .filter(|&n| n <= MAX_ELEMS)
            .ok_or("CTM array too large")?;
        let unpacked = elems * 4;
        let packed_size = self.u32()? as usize;
        let props = self.bytes(5)?.to_vec();
        let compressed = self.bytes(packed_size)?;

        // OpenCTM stores raw LZMA (5 props, size out-of-band); prepend the
        // 8-byte uncompressed size to form a .lzma stream for lzma-rs.
        let mut stream = Vec::with_capacity(13 + compressed.len());
        stream.extend_from_slice(&props);
        stream.extend_from_slice(&(unpacked as u64).to_le_bytes());
        stream.extend_from_slice(compressed);

        let mut out = Vec::new();
        lzma_rs::lzma_decompress(&mut Cursor::new(stream), &mut out)
            .map_err(|e| format!("LZMA decode failed: {e:?}"))?;
        if out.len() != unpacked {
            return Err(format!("LZMA size mismatch: {} != {}", out.len(), unpacked));
        }

        // Decompressed data is 4 big-endian byte planes (MSB first), each
        // count*size bytes, indexed component-major (k*count + i).
        let plane = count * size;
        let mut words = vec![0u32; count * size];
        for i in 0..count {
            for k in 0..size {
                let pos = k * count + i;
                words[i * size + k] = ((out[pos] as u32) << 24)
                    | ((out[plane + pos] as u32) << 16)
                    | ((out[2 * plane + pos] as u32) << 8)
                    | (out[3 * plane + pos] as u32);
            }
        }
        Ok(words)
    }

    fn packed_floats(&mut self, count: usize, size: usize) -> Result<Vec<f32>, String> {
        Ok(self
            .packed_words(count, size)?
            .into_iter()
            .map(f32::from_bits)
            .collect())
    }
}

// Inverse of the MG1 index delta transform; wraps like the C unsigned `+=`.
fn restore_indices(idx: &mut [u32], tri_count: usize) {
    for i in 0..tri_count {
        if i >= 1 {
            idx[i * 3] = idx[i * 3].wrapping_add(idx[(i - 1) * 3]);
        }
        idx[i * 3 + 2] = idx[i * 3 + 2].wrapping_add(idx[i * 3]);
        if i >= 1 && idx[i * 3] == idx[(i - 1) * 3] {
            idx[i * 3 + 1] = idx[i * 3 + 1].wrapping_add(idx[(i - 1) * 3 + 1]);
        } else {
            idx[i * 3 + 1] = idx[i * 3 + 1].wrapping_add(idx[i * 3]);
        }
    }
}

/// Decode an OpenCTM (MG1) mesh from memory.
pub fn decode(data: &[u8]) -> Result<CtmMesh, String> {
    let mut r = Reader { d: data, pos: 0 };
    r.magic(b"OCTM")?;
    let version = r.u32()?;
    if version != FORMAT_VERSION {
        return Err(format!("unsupported CTM version {version}"));
    }
    let method = r.bytes(4)?;
    if method != b"MG1\0" {
        return Err(format!(
            "unsupported CTM method {:?} (only MG1 is implemented)",
            String::from_utf8_lossy(method)
        ));
    }
    let vert_count = r.u32()? as usize;
    let tri_count = r.u32()? as usize;
    let uv_map_count = r.u32()?;
    let attr_map_count = r.u32()?;
    let flags = r.u32()?;
    r.skip_string()?;
    if vert_count == 0 || tri_count == 0 {
        return Err("CTM mesh has no vertices or triangles".into());
    }

    r.magic(b"INDX")?;
    let mut idx = r.packed_words(tri_count, 3)?;
    restore_indices(&mut idx, tri_count);

    r.magic(b"VERT")?;
    let vflat = r.packed_floats(vert_count * 3, 1)?;
    let positions = vflat.chunks_exact(3).map(|c| [c[0], c[1], c[2]]).collect();

    let normals = if flags & HAS_NORMALS_BIT != 0 {
        r.magic(b"NORM")?;
        let n = r.packed_floats(vert_count, 3)?;
        Some(n.chunks_exact(3).map(|c| [c[0], c[1], c[2]]).collect())
    } else {
        None
    };

    let mut uvs = None;
    for m in 0..uv_map_count {
        r.magic(b"TEXC")?;
        r.skip_string()?;
        r.skip_string()?;
        let uv = r.packed_floats(vert_count, 2)?;
        if m == 0 {
            uvs = Some(uv.chunks_exact(2).map(|c| [c[0], c[1]]).collect());
        }
    }

    for _ in 0..attr_map_count {
        r.magic(b"ATTR")?;
        r.skip_string()?;
        r.packed_floats(vert_count, 4)?;
    }

    Ok(CtmMesh {
        positions,
        normals,
        uvs,
        indices: idx,
    })
}

#[cfg(test)]
mod tests {
    const SAMPLE: &[u8] = include_bytes!("ctm_sample.ctm");

    #[test]
    fn decodes_reference_mesh() {
        let m = super::decode(SAMPLE).unwrap();
        assert_eq!(m.positions.len(), 449);
        assert_eq!(m.indices.len(), 378 * 3);
        assert_eq!(m.uvs.as_ref().unwrap().len(), 449);
        assert!(m.normals.is_none());
        assert!(m.indices.iter().all(|&i| (i as usize) < 449));
        let sum: f32 = m.positions.iter().flatten().sum();
        assert!((sum - 14352.117).abs() < 0.05);
    }

    #[test]
    fn rejects_oversized_count() {
        let mut h = Vec::new();
        h.extend_from_slice(b"OCTM");
        h.extend_from_slice(&5u32.to_le_bytes());
        h.extend_from_slice(b"MG1\0");
        h.extend_from_slice(&1u32.to_le_bytes());
        h.extend_from_slice(&u32::MAX.to_le_bytes());
        h.extend_from_slice(&0u32.to_le_bytes());
        h.extend_from_slice(&0u32.to_le_bytes());
        h.extend_from_slice(&0u32.to_le_bytes());
        h.extend_from_slice(&0u32.to_le_bytes());
        h.extend_from_slice(b"INDX");
        assert!(super::decode(&h).is_err());
    }
}