mxmxmx2tiles 0.1.0

//! meshopt-compressed glTF (EXT_meshopt_compression). Buffer 0 is the GLB BIN:
//! the compressed streams followed by the raw JPEGs. Buffer 1 is the fallback
//! (allocate-on-load) holding the decompressed layout. `quantize` adds
//! KHR_mesh_quantization (uint16 positions, ~mm loss, smaller); otherwise
//! positions/UVs stay float32 and the geometry is lossless.

use crate::glb::Prim;
use serde_json::{json, Value};

const FLOAT: u32 = 5126;
const USHORT: u32 = 5123;
const UINT: u32 = 5125;

struct Builder {
    comp: Vec<u8>,
    fb_len: usize,
    views: Vec<Value>,
    accessors: Vec<Value>,
}

impl Builder {
    fn align4(v: usize) -> usize {
        (v + 3) & !3
    }

    fn add(&mut self, decompressed_len: usize, stride: usize, count: usize, mode: &str, compressed: &[u8]) -> usize {
        self.fb_len = Self::align4(self.fb_len);
        let fb_off = self.fb_len;
        self.fb_len += decompressed_len;
        while self.comp.len() % 4 != 0 {
            self.comp.push(0);
        }
        let c_off = self.comp.len();
        self.comp.extend_from_slice(compressed);

        let mut bv = json!({
            "buffer": 1,
            "byteOffset": fb_off,
            "byteLength": decompressed_len,
            "extensions": { "EXT_meshopt_compression": {
                "buffer": 0,
                "byteOffset": c_off,
                "byteLength": compressed.len(),
                "byteStride": stride,
                "count": count,
                "mode": mode,
            }},
        });
        if mode == "ATTRIBUTES" {
            bv["byteStride"] = json!(stride);
        }
        self.views.push(bv);
        self.views.len() - 1
    }
}

// meshopt's index codec asserts (aborts) on a non-multiple-of-3 index count.
fn encodable(p: &Prim) -> bool {
    !p.positions.is_empty() && !p.indices.is_empty() && p.indices.len() % 3 == 0
}

fn enc<T>(rows: &[T]) -> Result<Vec<u8>, String> {
    meshopt::encode_vertex_buffer(rows).map_err(|e| format!("encode attribute: {e:?}"))
}

pub fn write(prims: &[Prim], quantize: bool) -> Result<Vec<u8>, String> {
    let usable: Vec<&Prim> = prims.iter().filter(|p| encodable(p)).collect();
    if usable.is_empty() {
        return Ok(crate::glb::write(prims));
    }

    let (offset, inv, scale) = if quantize {
        let mut all = Vec::new();
        for p in &usable {
            for v in &p.positions {
                all.extend_from_slice(v);
            }
        }
        let (o, s) = meshopt::calc_pos_offset_and_scale(&all);
        (o, if s > 0.0 { 65535.0 / s } else { 0.0 }, s)
    } else {
        ([0.0; 3], 0.0, 0.0)
    };

    let mut b = Builder { comp: Vec::new(), fb_len: 0, views: Vec::new(), accessors: Vec::new() };
    let mut images: Vec<Value> = Vec::new();
    let mut textures: Vec<Value> = Vec::new();
    let mut materials: Vec<Value> = Vec::new();
    let mut gltf_prims: Vec<Value> = Vec::new();
    let mut image_bin: Vec<u8> = Vec::new();

    for p in &usable {
        let vcount = p.positions.len();

        let pos_acc = b.accessors.len();
        if quantize {
            let mut q: Vec<[u16; 4]> = Vec::with_capacity(vcount);
            let mut lo = [u16::MAX; 3];
            let mut hi = [u16::MIN; 3];
            for v in &p.positions {
                let c = |i: usize| (((v[i] - offset[i]) * inv).round()).clamp(0.0, 65535.0) as u16;
                let w = [c(0), c(1), c(2)];
                for k in 0..3 {
                    lo[k] = lo[k].min(w[k]);
                    hi[k] = hi[k].max(w[k]);
                }
                q.push([w[0], w[1], w[2], 0]);
            }
            let view = b.add(q.len() * 8, 8, vcount, "ATTRIBUTES", &enc(&q)?);
            b.accessors.push(json!({
                "bufferView": view, "componentType": USHORT, "count": vcount,
                "type": "VEC3", "min": lo, "max": hi,
            }));
        } else {
            let mut lo = [f32::INFINITY; 3];
            let mut hi = [f32::NEG_INFINITY; 3];
            for v in &p.positions {
                for k in 0..3 {
                    lo[k] = lo[k].min(v[k]);
                    hi[k] = hi[k].max(v[k]);
                }
            }
            let view = b.add(vcount * 12, 12, vcount, "ATTRIBUTES", &enc(&p.positions)?);
            b.accessors.push(json!({
                "bufferView": view, "componentType": FLOAT, "count": vcount,
                "type": "VEC3", "min": lo, "max": hi,
            }));
        }

        let mut attributes = json!({ "POSITION": pos_acc });

        if let Some(uvs) = &p.uvs {
            let uv_acc = b.accessors.len();
            if quantize {
                let q: Vec<[u16; 2]> = uvs
                    .iter()
                    .map(|t| [(t[0].clamp(0.0, 1.0) * 65535.0).round() as u16, (t[1].clamp(0.0, 1.0) * 65535.0).round() as u16])
                    .collect();
                let view = b.add(q.len() * 4, 4, vcount, "ATTRIBUTES", &enc(&q)?);
                b.accessors.push(json!({
                    "bufferView": view, "componentType": USHORT, "count": vcount,
                    "type": "VEC2", "normalized": true,
                }));
            } else {
                let view = b.add(vcount * 8, 8, vcount, "ATTRIBUTES", &enc(uvs)?);
                b.accessors.push(json!({
                    "bufferView": view, "componentType": FLOAT, "count": vcount, "type": "VEC2",
                }));
            }
            attributes["TEXCOORD_0"] = json!(uv_acc);
        }

        let comp = meshopt::encode_index_buffer(&p.indices, vcount).map_err(|e| format!("encode indices: {e:?}"))?;
        let iv = b.add(p.indices.len() * 4, 4, p.indices.len(), "TRIANGLES", &comp);
        let idx_acc = b.accessors.len();
        b.accessors.push(json!({
            "bufferView": iv, "componentType": UINT, "count": p.indices.len(), "type": "SCALAR",
        }));

        let mut prim = json!({ "attributes": attributes, "indices": idx_acc });

        if let Some(jpeg) = &p.jpeg {
            while image_bin.len() % 4 != 0 {
                image_bin.push(0);
            }
            let img_off = image_bin.len();
            image_bin.extend_from_slice(jpeg);
            let view_idx = b.views.len();
            b.views.push(json!({ "buffer": 0, "byteOffset": img_off, "byteLength": jpeg.len(), "__image": true }));
            let img = images.len();
            images.push(json!({ "bufferView": view_idx, "mimeType": "image/jpeg" }));
            let tex = textures.len();
            textures.push(json!({ "source": img, "sampler": 0 }));
            let mat = materials.len();
            materials.push(json!({
                "pbrMetallicRoughness": { "baseColorTexture": { "index": tex }, "metallicFactor": 0.0, "roughnessFactor": 1.0 },
                "extensions": { "KHR_materials_unlit": {} },
                "doubleSided": true,
            }));
            prim["material"] = json!(mat);
        }
        gltf_prims.push(prim);
    }

    let img_base = Builder::align4(b.comp.len());
    for v in b.views.iter_mut() {
        if v.get("__image").is_some() {
            let off = v["byteOffset"].as_u64().unwrap() as usize + img_base;
            v["byteOffset"] = json!(off);
            v.as_object_mut().unwrap().remove("__image");
        }
    }
    while b.comp.len() < img_base {
        b.comp.push(0);
    }
    b.comp.extend_from_slice(&image_bin);

    let node = if quantize {
        let s = scale / 65535.0;
        json!({ "mesh": 0, "matrix": [s, 0.0, 0.0, 0.0, 0.0, s, 0.0, 0.0, 0.0, 0.0, s, 0.0, offset[0], offset[1], offset[2], 1.0] })
    } else {
        json!({ "mesh": 0 })
    };

    let mut ext_used = vec!["EXT_meshopt_compression"];
    let mut ext_req = vec!["EXT_meshopt_compression"];
    if quantize {
        ext_used.push("KHR_mesh_quantization");
        ext_req.push("KHR_mesh_quantization");
    }
    if !images.is_empty() {
        ext_used.push("KHR_materials_unlit");
    }

    let mut gltf = json!({
        "asset": { "version": "2.0", "generator": "mxmxmx2tiles" },
        "extensionsUsed": ext_used,
        "extensionsRequired": ext_req,
        "scene": 0,
        "scenes": [ { "nodes": [0] } ],
        "nodes": [ node ],
        "meshes": [ { "primitives": gltf_prims } ],
        "accessors": b.accessors,
        "bufferViews": b.views,
        "buffers": [
            { "byteLength": b.comp.len() },
            { "byteLength": b.fb_len, "extensions": { "EXT_meshopt_compression": { "fallback": true } } },
        ],
    });
    if !images.is_empty() {
        gltf["images"] = json!(images);
        gltf["textures"] = json!(textures);
        gltf["materials"] = json!(materials);
        // CLAMP_TO_EDGE, LINEAR_MIPMAP_LINEAR, LINEAR
        gltf["samplers"] = json!([ { "wrapS": 33071, "wrapT": 33071, "minFilter": 9987, "magFilter": 9729 } ]);
    }

    Ok(crate::glb::pack_glb(&gltf, &b.comp))
}

#[cfg(test)]
mod tests {
    use crate::glb::Prim;

    fn sample() -> Prim {
        Prim {
            positions: (0..500).map(|i| [i as f32 * 0.1, (i % 7) as f32 * 2.0, (i % 13) as f32 * -1.5]).collect(),
            normals: None,
            uvs: Some(vec![[0.25, 0.75]; 500]),
            indices: (0..(498 * 3)).map(|i| (i % 500) as u32).collect(),
            jpeg: None,
        }
    }

    #[test]
    fn lossless_roundtrip_is_exact() {
        let p = sample();
        let glb = super::write(&[p.clone()], false).unwrap();
        let jl = u32::from_le_bytes(glb[12..16].try_into().unwrap()) as usize;
        let g: serde_json::Value = serde_json::from_slice(&glb[20..20 + jl]).unwrap();
        let bin = &glb[20 + jl + 8..];
        let acc = &g["accessors"][0];
        assert_eq!(acc["componentType"], 5126);
        assert!(g["nodes"][0].get("matrix").is_none());
        let bv = &g["bufferViews"][acc["bufferView"].as_u64().unwrap() as usize];
        let e = &bv["extensions"]["EXT_meshopt_compression"];
        let o = e["byteOffset"].as_u64().unwrap() as usize;
        let l = e["byteLength"].as_u64().unwrap() as usize;
        let dec: Vec<[f32; 3]> = meshopt::decode_vertex_buffer(&bin[o..o + l], p.positions.len()).unwrap();
        assert_eq!(dec, p.positions);
    }

    #[test]
    fn quantized_roundtrip_within_step() {
        let p = sample();
        let glb = super::write(&[p.clone()], true).unwrap();
        let jl = u32::from_le_bytes(glb[12..16].try_into().unwrap()) as usize;
        let g: serde_json::Value = serde_json::from_slice(&glb[20..20 + jl]).unwrap();
        let bin = &glb[20 + jl + 8..];
        let m = g["nodes"][0]["matrix"].as_array().unwrap();
        let s = m[0].as_f64().unwrap() as f32;
        let off = [m[12].as_f64().unwrap() as f32, m[13].as_f64().unwrap() as f32, m[14].as_f64().unwrap() as f32];
        let acc = &g["accessors"][0];
        let bv = &g["bufferViews"][acc["bufferView"].as_u64().unwrap() as usize];
        let e = &bv["extensions"]["EXT_meshopt_compression"];
        let o = e["byteOffset"].as_u64().unwrap() as usize;
        let l = e["byteLength"].as_u64().unwrap() as usize;
        let dec: Vec<[u16; 4]> = meshopt::decode_vertex_buffer(&bin[o..o + l], p.positions.len()).unwrap();
        for (orig, q) in p.positions.iter().zip(dec.iter()) {
            for k in 0..3 {
                assert!((q[k] as f32 * s + off[k] - orig[k]).abs() <= s * 2.0);
            }
        }
    }
}