cyborg/src/pass/mesh.rs

use super::*;
use crate::shader::{ShaderHandle, ShaderStore};
use crate::storage::mesh::*;
use crate::storage::GpuVec;
use crate::viewport::ViewportInfo;
use crate::RenderLayouts;
use parking_lot::RwLock;

#[derive(Clone)]
pub struct ShaderInfo {
    pub store: Arc<ShaderStore>,
    pub forward: ShaderHandle,
    pub skinning: ShaderHandle,
}

#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct Vertex {
    pub position: [f32; 3],
    pub tan_frame: u32,
}

impl Attribute for Vertex {
    fn get_usages() -> wgpu::BufferUsages {
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::VERTEX
    }
}

const VERTEX_ATTRS: &[wgpu::VertexAttribute] =
    &wgpu::vertex_attr_array![0 => Float32x3, 1 => Uint32];

impl Vertex {
    pub fn desc() -> wgpu::VertexBufferLayout<'static> {
        wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<Self>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: VERTEX_ATTRS,
        }
    }
}

pub type Index = u32;

#[derive(Clone)]
pub struct Attributes {
    pub vertex: AttrId,
    pub index: AttrId,
}

impl Attributes {
    pub fn new(attr_store: impl AsRef<AttrStore>) -> Self {
        Self {
            vertex: attr_store.as_ref().get_type::<Vertex>(),
            index: attr_store.as_ref().add(AttrInfo {
                layout: AttrLayout {
                    size: std::mem::size_of::<Index>(),
                },
                usages: wgpu::BufferUsages::INDEX,
                default_pool_size: 1_000_000,
            }),
        }
    }
}

#[derive(Clone, Debug)]
pub struct TransformedMesh {
    pub mesh: MeshHandle,
    pub transform: glam::Mat4,
}

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct SkinningUniform {
    transform: [[f32; 4]; 4],
    src_offset: u32,
    dst_offset: u32,
    count: u32,
    _padding: u32,
}

#[derive(Debug)]
struct MeshCommand {
    vertex_offset: usize,
    vertex_count: usize,
    skinned_offset: usize,
    index_offset: usize,
    index_count: usize,
    skinning_index: usize,
}

#[derive(Debug)]
struct MeshGroupCommands {
    binding_indices: MeshLayoutBindingIndices,
    bind_group: wgpu::BindGroup,
    meshes: Vec<MeshCommand>,
}

pub struct FrameData {
    skinned_vertices: GpuVec<Vertex>,
    skinning_uniforms: GpuVec<SkinningUniform>,
    groups: Vec<MeshGroupCommands>,
}

pub struct MeshPass {
    device: Arc<wgpu::Device>,
    layouts: Arc<RenderLayouts>,
    attr_store: Arc<AttrStore>,
    shader_info: ShaderInfo,
    mesh_pool: Arc<MeshPool>,
    attributes: Attributes,
    mesh_layout_id: MeshLayoutId,
    skinning_bind_group_layout: wgpu::BindGroupLayout,
    skinning_pipeline: wgpu::ComputePipeline,
    depth_pipeline: wgpu::RenderPipeline,
    opaque_pipeline: wgpu::RenderPipeline,
    target_info: ViewportInfo,
    transformed_meshes: RwLock<Vec<TransformedMesh>>,
}

impl MeshPass {
    pub fn new(
        device: Arc<wgpu::Device>,
        layouts: Arc<RenderLayouts>,
        target_info: ViewportInfo,
        shader_info: ShaderInfo,
    ) -> Self {
        let attr_store = AttrStore::new();
        let mesh_pool = MeshPool::new(device.clone(), attr_store.to_owned());
        let attributes = Attributes::new(&attr_store);

        let mut mesh_layout = MeshLayoutDesc::new();
        mesh_layout.insert(attributes.vertex, ());
        mesh_layout.insert(attributes.index, ());
        let mesh_layout_id = mesh_pool.add_layout(mesh_layout).unwrap();

        let render_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("MeshPass Pipeline Layout"),
                bind_group_layouts: &[&layouts.bind_viewport],
                push_constant_ranges: &[],
            });

        let shader = shader_info.store.get(&shader_info.forward).unwrap();

        let targets = &[Some(wgpu::ColorTargetState {
            format: target_info.output_format,
            blend: Some(wgpu::BlendState::REPLACE),
            write_mask: wgpu::ColorWrites::ALL,
        })];

        let mut pipeline_desc = wgpu::RenderPipelineDescriptor {
            label: Some("Opaque MeshPass Pipeline"),
            layout: Some(&render_pipeline_layout),
            vertex: wgpu::VertexState {
                module: shader.as_ref(),
                entry_point: "vs_main",
                buffers: &[Vertex::desc()],
            },
            fragment: Some(wgpu::FragmentState {
                module: shader.as_ref(),
                entry_point: "fs_main",
                targets,
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: None, // Some(wgpu::Face::Back),
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: Some(wgpu::DepthStencilState {
                format: target_info.depth_format,
                depth_write_enabled: true,
                depth_compare: wgpu::CompareFunction::Less,
                stencil: Default::default(),
                bias: Default::default(),
            }),
            multisample: wgpu::MultisampleState {
                count: 1,
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            multiview: None,
        };

        let depth_pipeline = device.create_render_pipeline(&pipeline_desc);

        pipeline_desc.depth_stencil = Some(wgpu::DepthStencilState {
            format: target_info.depth_format,
            depth_write_enabled: false,
            depth_compare: wgpu::CompareFunction::Equal,
            stencil: Default::default(),
            bias: Default::default(),
        });

        let opaque_pipeline = device.create_render_pipeline(&pipeline_desc);

        drop(shader);

        let skinning_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("Skinning Bind Group Layout"),
                entries: &[
                    wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: true,
                            min_binding_size: None, // TODO ???
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 1,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: false },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 2,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    },
                ],
            });

        let skinning_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("Skinning Pipeline Layout"),
                bind_group_layouts: &[&skinning_bind_group_layout],
                push_constant_ranges: &[],
            });

        let skinning_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("Skinning Pipeline"),
            layout: Some(&skinning_pipeline_layout),
            module: shader_info
                .store
                .get(&shader_info.skinning)
                .unwrap()
                .as_ref(),
            entry_point: "cs_main",
        });

        Self {
            device,
            layouts,
            attr_store,
            shader_info,
            mesh_pool,
            attributes,
            mesh_layout_id,
            skinning_bind_group_layout,
            skinning_pipeline,
            depth_pipeline,
            opaque_pipeline,
            target_info,
            transformed_meshes: Default::default(),
        }
    }

    pub fn get_mesh_pool(&self) -> &MeshPool {
        &self.mesh_pool
    }

    pub fn get_attributes(&self) -> &Attributes {
        &self.attributes
    }

    pub fn add_transformed_meshes(&self, meshes: &[TransformedMesh]) {
        self.transformed_meshes.write().extend_from_slice(meshes);
    }
}

impl RenderPass for MeshPass {
    type FrameData = FrameData;

    fn create_frame_data(&self) -> FrameData {
        FrameData {
            skinned_vertices: GpuVec::new(
                self.device.clone(),
                Vertex::get_usages(),
                1024 * 128,
                Some("Skinned Vertices".to_string()),
                false,
            ),
            skinning_uniforms: GpuVec::new(
                self.device.clone(),
                wgpu::BufferUsages::STORAGE,
                1024 * 128,
                Some("Skinning Uniforms".to_string()),
                true,
            ),
            groups: Default::default(),
        }
    }

    fn begin_frame(&self, data: &mut FrameData, phases: &mut Vec<Phase>, queue: &wgpu::Queue) {
        phases.push(Phase::Upload);
        phases.push(Phase::Skinning);
        phases.push(Phase::Depth);
        phases.push(Phase::Opaque);
        phases.push(Phase::Transparent);

        data.groups.clear();
        data.skinning_uniforms.clear();

        let mut transformed_meshes_lock = self.transformed_meshes.write();
        let mesh_bindings = self
            .mesh_pool
            .iter_meshes(self.mesh_layout_id, transformed_meshes_lock.iter(), |i| {
                &i.mesh
            })
            .unwrap();

        let mut skinned_cursor = 0;

        for MeshLayoutInstances {
            bindings,
            instances,
        } in mesh_bindings.iter()
        {
            let pools = self.mesh_pool.get_bindings(bindings.clone());
            let vertices_pool = pools.get(self.attributes.vertex).unwrap();

            // TODO defer bind group creation into separate Vec after GpuVecs have been written
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &self.skinning_bind_group_layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                            buffer: data.skinning_uniforms.as_ref(),
                            offset: 0,
                            // TODO ugly!
                            size: Some(
                                std::num::NonZeroU64::new(
                                    std::mem::size_of::<SkinningUniform>() as u64
                                )
                                .unwrap(),
                            ),
                        }),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                            buffer: data.skinned_vertices.as_ref(),
                            offset: 0,
                            size: None,
                        }),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                            buffer: vertices_pool.get_buffer(),
                            offset: 0,
                            size: None,
                        }),
                    },
                ],
            });

            let mut group = MeshGroupCommands {
                binding_indices: bindings.clone(),
                meshes: Default::default(),
                bind_group,
            };

            for (mesh, infos) in instances {
                let vertices = infos
                    .iter()
                    .find(|i| i.0 == self.attributes.vertex)
                    .unwrap()
                    .1;
                let indices = infos
                    .iter()
                    .find(|i| i.0 == self.attributes.index)
                    .unwrap()
                    .1;

                group.meshes.push(MeshCommand {
                    vertex_offset: vertices.offset,
                    vertex_count: vertices.count,
                    skinned_offset: skinned_cursor,
                    index_offset: indices.offset,
                    index_count: indices.count,
                    skinning_index: data.skinning_uniforms.len(),
                });

                data.skinning_uniforms.push(SkinningUniform {
                    transform: mesh.transform.to_cols_array_2d(),
                    src_offset: vertices.offset as u32,
                    dst_offset: skinned_cursor as u32,
                    count: vertices.count as u32,
                    _padding: 0,
                });

                skinned_cursor += vertices.count;
            }

            data.groups.push(group);
        }

        data.skinned_vertices.reserve(skinned_cursor);
        data.skinned_vertices.write(&queue);
        data.skinning_uniforms.write(&queue);

        transformed_meshes_lock.clear();
    }

    fn record_commands(&self, data: PhaseData<&FrameData>, cmds: &mut wgpu::CommandEncoder) {
        match data.phase {
            Phase::Upload => self.mesh_pool.flush(cmds),
            _ => {}
        }
    }

    fn record_compute<'a>(
        &'a self,
        data: PhaseData<&'a FrameData>,
        cmds: &mut wgpu::ComputePass<'a>,
    ) {
        cmds.set_pipeline(&self.skinning_pipeline);
        for group in data.frame_data.groups.iter() {
            for mesh in group.meshes.iter() {
                let ubo_offset = data
                    .frame_data
                    .skinning_uniforms
                    .buf_offset(mesh.skinning_index);
                cmds.set_bind_group(0, &group.bind_group, &[ubo_offset as u32]);

                // TODO use div_ceil instead
                let workgroup_num = if mesh.vertex_count % 64 == 0 {
                    mesh.vertex_count / 64
                } else {
                    mesh.vertex_count / 64 + 1
                };

                cmds.dispatch_workgroups(workgroup_num as u32, 1, 1);
            }
        }
    }

    fn record_render(&self, data: PhaseData<&FrameData>) -> Option<wgpu::RenderBundle> {
        let pipeline = match data.phase {
            Phase::Depth => &self.depth_pipeline,
            Phase::Opaque => &self.opaque_pipeline,
            _ => return None,
        };

        let mut cmds =
            self.device
                .create_render_bundle_encoder(&wgpu::RenderBundleEncoderDescriptor {
                    label: Some("Opaque Pass Render Bundle"),
                    color_formats: &[Some(self.target_info.output_format)],
                    depth_stencil: Some(wgpu::RenderBundleDepthStencil {
                        format: self.target_info.depth_format,
                        depth_read_only: false, // TODO optimize?
                        stencil_read_only: true,
                    }),
                    sample_count: 1,
                    multiview: None,
                });

        // yikes
        let mesh_bindings: Vec<_> = data
            .frame_data
            .groups
            .iter()
            .map(
                |MeshGroupCommands {
                     binding_indices,
                     meshes,
                     ..
                 }| (self.mesh_pool.get_bindings(binding_indices.clone()), meshes),
            )
            .collect();

        cmds.set_pipeline(pipeline);
        cmds.set_bind_group(0, data.bind_viewport, &[]);

        for (bindings, meshes) in mesh_bindings.iter() {
            let indices_pool = bindings.get(self.attributes.index).unwrap();

            cmds.set_vertex_buffer(0, data.frame_data.skinned_vertices.as_ref().slice(..));
            cmds.set_index_buffer(
                indices_pool.get_buffer().slice(..),
                wgpu::IndexFormat::Uint32,
            );

            for mesh in meshes.iter() {
                let is_start = mesh.index_offset as u32;
                let is_end = is_start + mesh.index_count as u32;
                let vs_offset = mesh.skinned_offset as i32;
                cmds.draw_indexed(is_start..is_end, vs_offset, 0..1);
            }
        }

        Some(cmds.finish(&wgpu::RenderBundleDescriptor::default()))
    }
}