nightshade-renderer 0.57.0

use lz4_flex::frame::{FrameDecoder, FrameEncoder};
use std::io::{Read, Write};
use std::sync::Arc;

pub const MESHLET_MESH_ASSET_MAGIC: u64 = 1_717_551_717_668;
pub const MESHLET_MESH_ASSET_VERSION: u64 = 3;

/// A mesh pre-processed into a hierarchy of small triangle clusters (meshlets)
/// plus a BVH8 level-of-detail DAG, ready for GPU-driven virtualized-geometry
/// rendering. Produced offline by `MeshletMesh::from_mesh` behind the
/// `meshlet_processor` feature.
#[derive(Clone)]
pub struct MeshletMesh {
    pub vertex_positions: Arc<[u32]>,
    pub vertex_normals: Arc<[u32]>,
    pub vertex_uvs: Arc<[[f32; 2]]>,
    pub indices: Arc<[u8]>,
    pub bvh: Arc<[BvhNode]>,
    pub meshlets: Arc<[Meshlet]>,
    pub meshlet_cull_data: Arc<[MeshletCullData]>,
    pub aabb: MeshletAabb,
    pub bvh_depth: u32,
}

#[repr(C)]
#[derive(Copy, Clone, Default, bytemuck::Pod, bytemuck::Zeroable)]
pub struct BvhNode {
    pub aabbs: [MeshletAabbErrorOffset; 8],
    pub lod_bounds: [MeshletBoundingSphere; 8],
    pub child_counts: [u8; 8],
    pub _padding: [u32; 2],
}

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
pub struct Meshlet {
    pub start_vertex_position_bit: u32,
    pub start_vertex_attribute_id: u32,
    pub start_index_id: u32,
    pub vertex_count_minus_one: u8,
    pub triangle_count: u8,
    pub padding: u16,
    pub bits_per_vertex_position_channel_x: u8,
    pub bits_per_vertex_position_channel_y: u8,
    pub bits_per_vertex_position_channel_z: u8,
    pub vertex_position_quantization_factor: u8,
    pub min_vertex_position_channel_x: f32,
    pub min_vertex_position_channel_y: f32,
    pub min_vertex_position_channel_z: f32,
}

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
pub struct MeshletCullData {
    pub aabb: MeshletAabbErrorOffset,
    pub lod_group_sphere: MeshletBoundingSphere,
}

#[repr(C)]
#[derive(Copy, Clone, Default, bytemuck::Pod, bytemuck::Zeroable)]
pub struct MeshletAabb {
    pub center: [f32; 3],
    pub _pad0: f32,
    pub half_extent: [f32; 3],
    pub _pad1: f32,
}

#[repr(C)]
#[derive(Copy, Clone, Default, bytemuck::Pod, bytemuck::Zeroable)]
pub struct MeshletAabbErrorOffset {
    pub center: [f32; 3],
    pub error: f32,
    pub half_extent: [f32; 3],
    pub child_offset: u32,
}

#[repr(C)]
#[derive(Copy, Clone, Default, bytemuck::Pod, bytemuck::Zeroable)]
pub struct MeshletBoundingSphere {
    pub center: [f32; 3],
    pub radius: f32,
}

#[derive(Debug, thiserror::Error)]
pub enum MeshletMeshSaveOrLoadError {
    #[error("file was not a MeshletMesh asset")]
    WrongFileType,
    #[error("expected asset version {MESHLET_MESH_ASSET_VERSION} but found version {found}")]
    WrongVersion { found: u64 },
    #[error("failed to compress or decompress asset data")]
    CompressionOrDecompression(#[from] lz4_flex::frame::Error),
    #[error(transparent)]
    Io(#[from] std::io::Error),
}

impl MeshletMesh {
    /// Serializes the mesh to the versioned, LZ4-compressed on-disk format.
    pub fn write(&self, writer: &mut impl Write) -> Result<(), MeshletMeshSaveOrLoadError> {
        writer.write_all(&MESHLET_MESH_ASSET_MAGIC.to_le_bytes())?;
        writer.write_all(&MESHLET_MESH_ASSET_VERSION.to_le_bytes())?;
        writer.write_all(bytemuck::bytes_of(&self.aabb))?;
        writer.write_all(&self.bvh_depth.to_le_bytes())?;

        let mut writer = FrameEncoder::new(writer);
        write_slice(&self.vertex_positions, &mut writer)?;
        write_slice(&self.vertex_normals, &mut writer)?;
        write_slice(&self.vertex_uvs, &mut writer)?;
        write_slice(&self.indices, &mut writer)?;
        write_slice(&self.bvh, &mut writer)?;
        write_slice(&self.meshlets, &mut writer)?;
        write_slice(&self.meshlet_cull_data, &mut writer)?;
        writer.finish()?;
        Ok(())
    }

    /// Loads a mesh from the versioned, LZ4-compressed on-disk format.
    pub fn read(reader: &mut impl Read) -> Result<Self, MeshletMeshSaveOrLoadError> {
        let magic = read_u64(reader)?;
        if magic != MESHLET_MESH_ASSET_MAGIC {
            return Err(MeshletMeshSaveOrLoadError::WrongFileType);
        }
        let version = read_u64(reader)?;
        if version != MESHLET_MESH_ASSET_VERSION {
            return Err(MeshletMeshSaveOrLoadError::WrongVersion { found: version });
        }

        let mut aabb_bytes = [0u8; std::mem::size_of::<MeshletAabb>()];
        reader.read_exact(&mut aabb_bytes)?;
        let aabb = *bytemuck::from_bytes(&aabb_bytes);
        let bvh_depth = read_u32(reader)?;

        let reader = &mut FrameDecoder::new(reader);
        let vertex_positions = read_slice(reader)?;
        let vertex_normals = read_slice(reader)?;
        let vertex_uvs = read_slice(reader)?;
        let indices = read_slice(reader)?;
        let bvh = read_slice(reader)?;
        let meshlets = read_slice(reader)?;
        let meshlet_cull_data = read_slice(reader)?;

        Ok(Self {
            vertex_positions,
            vertex_normals,
            vertex_uvs,
            indices,
            bvh,
            meshlets,
            meshlet_cull_data,
            aabb,
            bvh_depth,
        })
    }
}

fn write_slice<T: bytemuck::Pod>(
    slice: &[T],
    writer: &mut impl Write,
) -> Result<(), MeshletMeshSaveOrLoadError> {
    writer.write_all(&(slice.len() as u64).to_le_bytes())?;
    writer.write_all(bytemuck::cast_slice(slice))?;
    Ok(())
}

fn read_slice<T: bytemuck::Pod>(
    reader: &mut impl Read,
) -> Result<Arc<[T]>, MeshletMeshSaveOrLoadError> {
    let count = read_u64(reader)? as usize;
    let mut bytes = vec![0u8; count * std::mem::size_of::<T>()];
    reader.read_exact(&mut bytes)?;
    Ok(bytemuck::cast_slice(&bytes).into())
}

fn read_u64(reader: &mut impl Read) -> Result<u64, std::io::Error> {
    let mut bytes = [0u8; 8];
    reader.read_exact(&mut bytes)?;
    Ok(u64::from_le_bytes(bytes))
}

fn read_u32(reader: &mut impl Read) -> Result<u32, std::io::Error> {
    let mut bytes = [0u8; 4];
    reader.read_exact(&mut bytes)?;
    Ok(u32::from_le_bytes(bytes))
}