cyborg/src/model.rs

use crate::handle::*;
use crate::mesh::*;
use crate::pool::{MaterialData, TextureData};
use crate::scene::MeshInstance;

pub trait OnLoad {
    fn load_mesh(&mut self, mesh_data: &MeshData) -> MeshHandle;
    fn load_texture(&mut self, texture_data: &TextureData) -> TextureHandle;
    fn load_material(&mut self, material_data: &MaterialData) -> MaterialHandle;
}

pub struct BasicMesh {
    pub mesh_handle: MeshHandle,
    pub material_handle: MaterialHandle,
}

impl BasicMesh {
    pub fn instantiate(&self, transform: glam::Mat4) -> MeshInstance {
        MeshInstance {
            mesh: self.mesh_handle,
            material: self.material_handle,
            transform,
        }
    }
}

pub struct ObjModel {
    mesh: BasicMesh,
}

impl ObjModel {
    pub fn load(mut on_load: impl OnLoad) -> Self {
        use tobj::*;

        let model_data = include_bytes!("assets/viking_room/model.obj").to_vec();
        let model_data = &mut model_data.as_slice();

        let load_options = LoadOptions {
            triangulate: true,
            single_index: true,
            ..Default::default()
        };

        let (models, _mats) =
            load_obj_buf(model_data, &load_options, |_| unimplemented!()).unwrap();

        let mut vertices = Vec::new();
        let mut indices = Vec::new();

        let m = models.first().unwrap();
        let index_base = vertices.len() as u32;
        for i in 0..m.mesh.positions.len() / 3 {
            let t = i * 3;
            vertices.push(Vertex {
                position: [
                    m.mesh.positions[t],
                    m.mesh.positions[t + 2],
                    -m.mesh.positions[t + 1],
                ],
                normal: [
                    m.mesh.normals[t],
                    m.mesh.normals[t + 2],
                    -m.mesh.normals[t + 1],
                ],
                tex_coords: [m.mesh.texcoords[i * 2], 1.0 - m.mesh.texcoords[i * 2 + 1]],
            });
        }

        indices.extend(m.mesh.indices.iter().map(|i| i + index_base));

        let mesh_data = MeshData { vertices, indices };
        let mesh_handle = on_load.load_mesh(&mesh_data);

        let albedo_data = include_bytes!("assets/viking_room/albedo.png");
        let albedo = image::load_from_memory(albedo_data).unwrap();

        use image::GenericImageView;
        let dimensions = albedo.dimensions();

        let albedo_rgb = albedo.as_rgb8().unwrap().to_vec();
        let mut albedo_rgba = Vec::<u8>::new();
        for rgb in albedo_rgb.chunks(3) {
            albedo_rgba.extend_from_slice(rgb);
            albedo_rgba.push(0xff);
        }

        let albedo_data = TextureData {
            width: dimensions.0,
            height: dimensions.1,
            data: albedo_rgba,
        };

        let albedo = on_load.load_texture(&albedo_data);

        let material_handle = on_load.load_material(&MaterialData {
            albedo,
            metallic_roughness: albedo, // TODO better placeholder MR texture
        });

        let mesh = BasicMesh {
            mesh_handle,
            material_handle,
        };
        Self { mesh }
    }

    pub fn draw(&self, meshes: &mut Vec<MeshInstance>, transform: glam::Mat4) {
        meshes.push(self.mesh.instantiate(transform));
    }
}

pub struct GltfModel {
    meshes: Vec<BasicMesh>,
}

impl GltfModel {
    pub fn load(on_load: impl OnLoad) -> Self {
        use gltf::*;

        let model_data = include_bytes!("assets/DamagedHelmet.glb");
        let model = Gltf::from_slice(model_data.as_slice()).unwrap();
        let loader = GltfLoader::load(&model, on_load);

        Self {
            meshes: loader.meshes,
        }
    }

    pub fn draw(&self, meshes: &mut Vec<MeshInstance>, transform: glam::Mat4) {
        for mesh in self.meshes.iter() {
            meshes.push(mesh.instantiate(transform));
        }
    }
}

pub struct GltfLoader<'a, T: OnLoad> {
    pub model: &'a gltf::Gltf,
    pub buffers: Vec<Vec<u8>>,
    pub on_load: T,
    pub meshes: Vec<BasicMesh>,
}

impl<'a, T: OnLoad> GltfLoader<'a, T> {
    pub fn load(model: &'a gltf::Gltf, on_load: T) -> Self {
        let buffers = Self::load_buffers(model);

        let mut loader = Self {
            model,
            buffers,
            on_load,
            meshes: Default::default(),
        };

        for scene in loader.model.scenes() {
            for node in scene.nodes() {
                loader.load_node(node);
            }
        }

        loader
    }

    pub fn load_buffers(model: &gltf::Gltf) -> Vec<Vec<u8>> {
        let mut buffer_data = Vec::<Vec<u8>>::new();
        for buffer in model.buffers() {
            match buffer.source() {
                gltf::buffer::Source::Bin => {
                    buffer_data.push(model.blob.as_deref().unwrap().into());
                }
                _ => panic!("URI buffer sources are unsupported"),
            }
        }

        buffer_data
    }

    pub fn load_node(&mut self, node: gltf::Node) {
        if let Some(mesh) = node.mesh() {
            for primitive in mesh.primitives() {
                let mesh = self.load_primitive(primitive);
                self.meshes.push(mesh);
            }
        }

        for child in node.children() {
            self.load_node(child);
        }
    }

    pub fn load_primitive(&mut self, primitive: gltf::Primitive) -> BasicMesh {
        let material_handle = self.load_primitive_material(primitive.clone());
        let mesh_handle = self.load_primitive_mesh(primitive);
        BasicMesh {
            material_handle,
            mesh_handle,
        }
    }

    pub fn load_primitive_material(&mut self, primitive: gltf::Primitive) -> MaterialHandle {
        let pbr = primitive.material().pbr_metallic_roughness();

        let base_color = pbr.base_color_texture().unwrap().texture();
        let albedo = self.load_texture(base_color);

        let metallic_roughness = pbr.metallic_roughness_texture().unwrap().texture();
        let metallic_roughness = self.load_texture(metallic_roughness);

        self.on_load.load_material(&MaterialData {
            albedo,
            metallic_roughness,
        })
    }

    pub fn load_texture(&mut self, texture: gltf::Texture) -> TextureHandle {
        let source = texture.source().source();
        let view = if let gltf::image::Source::View { view, .. } = source {
            view
        } else {
            panic!("texture must be embedded");
        };

        let start = view.offset() as usize;
        let end = (view.offset() + view.length()) as usize;
        let src = &self.buffers[view.buffer().index()][start..end];

        use image::GenericImageView;
        let image = image::load_from_memory(src).unwrap();
        let dimensions = image.dimensions();

        let rgba: Vec<u8> = if let Some(rgba) = image.as_rgba8() {
            rgba.to_vec()
        } else {
            let rgb = image.as_rgb8().unwrap().to_vec();
            let mut rgba = Vec::<u8>::new();
            for rgb in rgb.chunks(3) {
                rgba.extend_from_slice(rgb);
                rgba.push(0xff);
            }
            rgba
        };

        self.on_load.load_texture(&TextureData {
            width: dimensions.0,
            height: dimensions.1,
            data: rgba,
        })
    }

    pub fn load_primitive_mesh(&mut self, primitive: gltf::Primitive) -> MeshHandle {
        use gltf::mesh::util::{ReadIndices, ReadTexCoords};

        if primitive.mode() != gltf::mesh::Mode::Triangles {
            panic!("glTF primitive must be triangle list");
        }

        let reader = primitive.reader(|buffer| Some(&self.buffers[buffer.index()]));
        let positions = reader.read_positions().unwrap();
        let mut normals = reader.read_normals().unwrap();
        let tex_coords = reader.read_tex_coords(0).unwrap();

        let mut tex_coords = if let ReadTexCoords::F32(tex_coords) = tex_coords {
            tex_coords
        } else {
            panic!("only f32 texture coordinates are supported")
        };

        let mut vertices = Vec::new();
        for position in positions {
            let normal = normals.next().unwrap();
            let tex_coords = tex_coords.next().unwrap();

            vertices.push(Vertex {
                position,
                normal,
                tex_coords,
            });
        }

        let indices = match reader.read_indices().unwrap() {
            ReadIndices::U32(indices) => indices.collect(),
            ReadIndices::U16(indices) => indices.map(|i| i as u32).collect(),
            ReadIndices::U8(indices) => indices.map(|i| i as u32).collect(),
        };

        self.on_load.load_mesh(&MeshData { vertices, indices })
    }
}