stylized 1-bit rendering

src/mesh.rs

@@ -1,5 +1,5 @@
 use bytemuck::{Pod, Zeroable};
-use glam::{Mat4, Vec3};
+use glam::{Mat4, Quat, Vec3};
 use std::path::Path;
 use std::rc::Rc;
 
@@ -42,6 +42,65 @@ impl Vertex
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct InstanceData
+{
+    pub position: Vec3,
+    pub rotation: Quat,
+    pub scale: Vec3,
+}
+
+impl InstanceData
+{
+    pub fn to_raw(&self) -> InstanceRaw
+    {
+        let model = Mat4::from_scale_rotation_translation(self.scale, self.rotation, self.position);
+        InstanceRaw {
+            model: model.to_cols_array_2d(),
+        }
+    }
+}
+
+#[repr(C)]
+#[derive(Clone, Copy, Pod, Zeroable)]
+pub struct InstanceRaw
+{
+    pub model: [[f32; 4]; 4],
+}
+
+impl InstanceRaw
+{
+    pub fn desc() -> wgpu::VertexBufferLayout<'static>
+    {
+        wgpu::VertexBufferLayout {
+            array_stride: std::mem::size_of::<InstanceRaw>() as wgpu::BufferAddress,
+            step_mode: wgpu::VertexStepMode::Instance,
+            attributes: &[
+                wgpu::VertexAttribute {
+                    offset: 0,
+                    shader_location: 3,
+                    format: wgpu::VertexFormat::Float32x4,
+                },
+                wgpu::VertexAttribute {
+                    offset: std::mem::size_of::<[f32; 4]>() as wgpu::BufferAddress,
+                    shader_location: 4,
+                    format: wgpu::VertexFormat::Float32x4,
+                },
+                wgpu::VertexAttribute {
+                    offset: (std::mem::size_of::<[f32; 4]>() * 2) as wgpu::BufferAddress,
+                    shader_location: 5,
+                    format: wgpu::VertexFormat::Float32x4,
+                },
+                wgpu::VertexAttribute {
+                    offset: (std::mem::size_of::<[f32; 4]>() * 3) as wgpu::BufferAddress,
+                    shader_location: 6,
+                    format: wgpu::VertexFormat::Float32x4,
+                },
+            ],
+        }
+    }
+}
+
 pub struct Mesh
 {
     pub vertex_buffer: wgpu::Buffer,
@@ -411,4 +470,271 @@ impl Mesh
     {
         crate::render::with_device(|device| Mesh::load_gltf_mesh(device, path))
     }
+
+    pub fn load_gltf_with_instances(
+        device: &wgpu::Device,
+        path: impl AsRef<Path>,
+    ) -> anyhow::Result<Vec<(Mesh, Vec<InstanceData>)>>
+    {
+        let path = path.as_ref();
+        let gltf_str = std::fs::read_to_string(path)?;
+        let gltf_json: serde_json::Value = serde_json::from_str(&gltf_str)?;
+
+        let (document, buffers, _images) = gltf::import(path)?;
+
+        let mut result = Vec::new();
+
+        let nodes = gltf_json["nodes"]
+            .as_array()
+            .ok_or_else(|| anyhow::anyhow!("Missing nodes array"))?;
+
+        for (node_index, json_node) in nodes.iter().enumerate()
+        {
+            let node = document
+                .nodes()
+                .nth(node_index)
+                .ok_or_else(|| anyhow::anyhow!("Node index mismatch"))?;
+
+            if let Some(mesh_data) = node.mesh()
+            {
+                let has_instancing = json_node
+                    .get("extensions")
+                    .and_then(|ext| ext.get("EXT_mesh_gpu_instancing"))
+                    .is_some();
+
+                if has_instancing
+                {
+                    let extensions = json_node.get("extensions").unwrap();
+                    let instancing_ext = extensions.get("EXT_mesh_gpu_instancing").unwrap();
+                        let mut mesh_vertices = Vec::new();
+                        let mut mesh_indices = Vec::new();
+
+                        for primitive in mesh_data.primitives()
+                        {
+                            let reader = primitive
+                                .reader(|buffer| buffers.get(buffer.index()).map(|data| &data[..]));
+
+                            let positions = reader
+                                .read_positions()
+                                .ok_or_else(|| anyhow::anyhow!("Missing position data"))?
+                                .collect::<Vec<[f32; 3]>>();
+
+                            let normals = reader
+                                .read_normals()
+                                .ok_or_else(|| anyhow::anyhow!("Missing normal data"))?
+                                .collect::<Vec<[f32; 3]>>();
+
+                            let uvs = reader
+                                .read_tex_coords(0)
+                                .map(|iter| iter.into_f32().collect::<Vec<[f32; 2]>>())
+                                .unwrap_or_else(|| vec![[0.0, 0.0]; positions.len()]);
+
+                            let base_index = mesh_vertices.len() as u32;
+
+                            for ((pos, normal), uv) in
+                                positions.iter().zip(normals.iter()).zip(uvs.iter())
+                            {
+                                mesh_vertices.push(Vertex {
+                                    position: *pos,
+                                    normal: *normal,
+                                    uv: *uv,
+                                });
+                            }
+
+                            if let Some(indices_reader) = reader.read_indices()
+                            {
+                                mesh_indices
+                                    .extend(indices_reader.into_u32().map(|i| i + base_index));
+                            }
+                        }
+
+                        let attributes = instancing_ext
+                            .get("attributes")
+                            .and_then(|v| v.as_object())
+                            .ok_or_else(|| anyhow::anyhow!("Missing attributes in EXT_mesh_gpu_instancing"))?;
+
+                        let translation_accessor_index = attributes
+                            .get("TRANSLATION")
+                            .and_then(|v| v.as_u64())
+                            .ok_or_else(|| anyhow::anyhow!("Missing TRANSLATION in instancing extension"))? as usize;
+
+                        let rotation_accessor_index = attributes
+                            .get("ROTATION")
+                            .and_then(|v| v.as_u64())
+                            .ok_or_else(|| anyhow::anyhow!("Missing ROTATION in instancing extension"))? as usize;
+
+                        let scale_accessor_index = attributes
+                            .get("SCALE")
+                            .and_then(|v| v.as_u64())
+                            .ok_or_else(|| anyhow::anyhow!("Missing SCALE in instancing extension"))? as usize;
+
+                        let translations = Self::read_vec3_accessor(
+                            &document,
+                            &buffers,
+                            translation_accessor_index,
+                        )?;
+                        let rotations =
+                            Self::read_quat_accessor(&document, &buffers, rotation_accessor_index)?;
+                        let scales =
+                            Self::read_vec3_accessor(&document, &buffers, scale_accessor_index)?;
+
+                        let instances: Vec<InstanceData> = translations
+                            .into_iter()
+                            .zip(rotations.into_iter())
+                            .zip(scales.into_iter())
+                            .map(|((position, rotation), scale)| InstanceData {
+                                position,
+                                rotation,
+                                scale,
+                            })
+                            .collect();
+
+                    let mesh = Mesh::new(device, &mesh_vertices, &mesh_indices);
+                    result.push((mesh, instances));
+                }
+                else
+                {
+                    let mut mesh_vertices = Vec::new();
+                    let mut mesh_indices = Vec::new();
+
+                    for primitive in mesh_data.primitives()
+                    {
+                        let reader = primitive
+                            .reader(|buffer| buffers.get(buffer.index()).map(|data| &data[..]));
+
+                        let positions = reader
+                            .read_positions()
+                            .ok_or_else(|| anyhow::anyhow!("Missing position data"))?
+                            .collect::<Vec<[f32; 3]>>();
+
+                        let normals = reader
+                            .read_normals()
+                            .ok_or_else(|| anyhow::anyhow!("Missing normal data"))?
+                            .collect::<Vec<[f32; 3]>>();
+
+                        let uvs = reader
+                            .read_tex_coords(0)
+                            .map(|iter| iter.into_f32().collect::<Vec<[f32; 2]>>())
+                            .unwrap_or_else(|| vec![[0.0, 0.0]; positions.len()]);
+
+                        let base_index = mesh_vertices.len() as u32;
+
+                        for ((pos, normal), uv) in
+                            positions.iter().zip(normals.iter()).zip(uvs.iter())
+                        {
+                            mesh_vertices.push(Vertex {
+                                position: *pos,
+                                normal: *normal,
+                                uv: *uv,
+                            });
+                        }
+
+                        if let Some(indices_reader) = reader.read_indices()
+                        {
+                            mesh_indices.extend(indices_reader.into_u32().map(|i| i + base_index));
+                        }
+                    }
+
+                    let mesh = Mesh::new(device, &mesh_vertices, &mesh_indices);
+                    result.push((mesh, Vec::new()));
+                }
+            }
+        }
+
+        Ok(result)
+    }
+
+    fn read_vec3_accessor(
+        document: &gltf::Document,
+        buffers: &[gltf::buffer::Data],
+        accessor_index: usize,
+    ) -> anyhow::Result<Vec<Vec3>>
+    {
+        let accessor = document
+            .accessors()
+            .nth(accessor_index)
+            .ok_or_else(|| anyhow::anyhow!("Invalid accessor index"))?;
+
+        let buffer_view = accessor.view().ok_or_else(|| anyhow::anyhow!("Missing buffer view"))?;
+        let buffer = &buffers[buffer_view.buffer().index()];
+        let start = buffer_view.offset() + accessor.offset();
+        let stride = buffer_view.stride().unwrap_or(12);
+
+        let mut result = Vec::new();
+        for i in 0..accessor.count()
+        {
+            let offset = start + i * stride;
+            let x = f32::from_le_bytes([
+                buffer[offset],
+                buffer[offset + 1],
+                buffer[offset + 2],
+                buffer[offset + 3],
+            ]);
+            let y = f32::from_le_bytes([
+                buffer[offset + 4],
+                buffer[offset + 5],
+                buffer[offset + 6],
+                buffer[offset + 7],
+            ]);
+            let z = f32::from_le_bytes([
+                buffer[offset + 8],
+                buffer[offset + 9],
+                buffer[offset + 10],
+                buffer[offset + 11],
+            ]);
+            result.push(Vec3::new(x, y, z));
+        }
+
+        Ok(result)
+    }
+
+    fn read_quat_accessor(
+        document: &gltf::Document,
+        buffers: &[gltf::buffer::Data],
+        accessor_index: usize,
+    ) -> anyhow::Result<Vec<Quat>>
+    {
+        let accessor = document
+            .accessors()
+            .nth(accessor_index)
+            .ok_or_else(|| anyhow::anyhow!("Invalid accessor index"))?;
+
+        let buffer_view = accessor.view().ok_or_else(|| anyhow::anyhow!("Missing buffer view"))?;
+        let buffer = &buffers[buffer_view.buffer().index()];
+        let start = buffer_view.offset() + accessor.offset();
+        let stride = buffer_view.stride().unwrap_or(16);
+
+        let mut result = Vec::new();
+        for i in 0..accessor.count()
+        {
+            let offset = start + i * stride;
+            let x = f32::from_le_bytes([
+                buffer[offset],
+                buffer[offset + 1],
+                buffer[offset + 2],
+                buffer[offset + 3],
+            ]);
+            let y = f32::from_le_bytes([
+                buffer[offset + 4],
+                buffer[offset + 5],
+                buffer[offset + 6],
+                buffer[offset + 7],
+            ]);
+            let z = f32::from_le_bytes([
+                buffer[offset + 8],
+                buffer[offset + 9],
+                buffer[offset + 10],
+                buffer[offset + 11],
+            ]);
+            let w = f32::from_le_bytes([
+                buffer[offset + 12],
+                buffer[offset + 13],
+                buffer[offset + 14],
+                buffer[offset + 15],
+            ]);
+            result.push(Quat::from_xyzw(x, y, z, w));
+        }
+
+        Ok(result)
+    }
 }