cyborg/src/pool.rs

use super::handle::{MaterialHandle, MeshHandle, TextureHandle};
use super::mesh::MeshData;
use slab::Slab;
use wgpu::util::DeviceExt;

pub struct MeshGroup {
    // TODO make these all private
    pub vertices: wgpu::Buffer,
    pub vertex_capacity: usize,
    pub indices: wgpu::Buffer,
    pub index_capacity: usize,
}

impl MeshGroup {
    fn new(device: &wgpu::Device, data: &MeshData) -> Self {
        let vertex_capacity = data.vertices.len();
        let vertices = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Vertex Buffer"),
            contents: bytemuck::cast_slice(&data.vertices),
            usage: wgpu::BufferUsages::VERTEX,
        });

        let index_capacity = data.indices.len();
        let indices = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Index Buffer"),
            contents: bytemuck::cast_slice(&data.indices),
            usage: wgpu::BufferUsages::INDEX,
        });

        Self {
            vertex_capacity,
            vertices,
            index_capacity,
            indices,
        }
    }
}

#[derive(Default)]
pub struct MeshPool {
    // TODO make this private
    pub groups: Slab<MeshGroup>,
}

impl MeshPool {
    pub fn allocate(&mut self, device: &wgpu::Device, data: &MeshData) -> MeshHandle {
        let group = MeshGroup::new(device, data);
        let group_id = self.groups.insert(group);
        let sub_id = 0;
        MeshHandle { group_id, sub_id }
    }

    pub fn get_group(&self, handle: &MeshHandle) -> Option<&MeshGroup> {
        self.groups.get(handle.group_id)
    }
}

pub struct TextureData {
    pub width: u32,
    pub height: u32,
    pub data: Vec<u8>,
}

pub struct Texture {
    texture: wgpu::Texture,
    texture_view: wgpu::TextureView,
}

impl Texture {
    pub fn new(device: &wgpu::Device, queue: &wgpu::Queue, data: &TextureData) -> Self {
        let size = wgpu::Extent3d {
            width: data.width,
            height: data.height,
            depth_or_array_layers: 1,
        };

        let texture = device.create_texture(&wgpu::TextureDescriptor {
            size,
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8UnormSrgb,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            label: None,
        });

        let texture_view = texture.create_view(&wgpu::TextureViewDescriptor::default());

        queue.write_texture(
            wgpu::ImageCopyTexture {
                texture: &texture,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            &data.data,
            wgpu::ImageDataLayout {
                offset: 0,
                bytes_per_row: std::num::NonZeroU32::new(4 * size.width),
                rows_per_image: std::num::NonZeroU32::new(size.height),
            },
            size,
        );

        Texture {
            texture,
            texture_view,
        }
    }
}

pub struct TexturePool {
    // TODO make this all private
    pub textures: Slab<Texture>,
    pub sampler: wgpu::Sampler,
}

impl TexturePool {
    pub fn new(device: &wgpu::Device) -> Self {
        let textures = Default::default();
        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            address_mode_u: wgpu::AddressMode::Repeat,
            address_mode_v: wgpu::AddressMode::Repeat,
            address_mode_w: wgpu::AddressMode::Repeat,
            mag_filter: wgpu::FilterMode::Linear,
            min_filter: wgpu::FilterMode::Nearest,
            mipmap_filter: wgpu::FilterMode::Nearest,
            ..Default::default()
        });

        Self { textures, sampler }
    }

    pub fn allocate(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        data: &TextureData,
    ) -> TextureHandle {
        let texture = Texture::new(device, queue, data);
        let id = self.textures.insert(texture);
        TextureHandle { id }
    }
}

pub struct MaterialData {
    pub albedo: TextureHandle,
    pub metallic_roughness: TextureHandle,
}

pub struct Material {
    pub bind_group: wgpu::BindGroup,
}

pub struct MaterialPool {
    pub materials: Slab<Material>,
    pub bind_group_layout: wgpu::BindGroupLayout,
}

impl MaterialPool {
    pub fn new(device: &wgpu::Device) -> Self {
        let texture_entry = wgpu::BindGroupLayoutEntry {
            binding: u32::MAX,
            visibility: wgpu::ShaderStages::FRAGMENT,
            ty: wgpu::BindingType::Texture {
                multisampled: false,
                view_dimension: wgpu::TextureViewDimension::D2,
                sample_type: wgpu::TextureSampleType::Float { filterable: true },
            },
            count: None,
        };

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    ..texture_entry
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    ..texture_entry
                },
            ],
            label: Some("Texture Bind Group Layout"),
        });

        Self {
            materials: Default::default(),
            bind_group_layout,
        }
    }

    pub fn allocate(
        &mut self,
        device: &wgpu::Device,
        texture_pool: &TexturePool,
        data: &MaterialData,
    ) -> MaterialHandle {
        let albedo_view = &texture_pool.textures[data.albedo.id].texture_view;
        let mr_view = &texture_pool.textures[data.metallic_roughness.id].texture_view;

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            layout: &self.bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::Sampler(&texture_pool.sampler),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(albedo_view),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::TextureView(mr_view),
                },
            ],
            label: None,
        });

        let material = Material { bind_group };
        let id = self.materials.insert(material);
        MaterialHandle { id }
    }
}