cyborg/src/pass/mesh.rs

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

pub struct ShaderInfo {
    pub store: Arc<ShaderStore>,
    pub forward: ShaderHandle,
}

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

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

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

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;

pub struct FrameData {}

pub struct MeshPass {
    device: Arc<wgpu::Device>,
    layouts: Arc<RenderLayouts>,
    attr_store: Arc<AttrStore>,
    shader_info: ShaderInfo,
    mesh_pool: Arc<MeshPool>,
    vertex_attr_id: AttrId,
    index_attr_id: AttrId,
    mesh_layout_id: MeshLayoutId,
    example_mesh: MeshHandle,
    depth_pipeline: wgpu::RenderPipeline,
    opaque_pipeline: wgpu::RenderPipeline,
    target_info: ViewportInfo,
}

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 vertex_attr_id = attr_store.get_type::<Vertex>();
        let index_attr_id = attr_store.add(AttrInfo {
            layout: AttrLayout {
                size: std::mem::size_of::<Index>(),
            },
            usages: wgpu::BufferUsages::INDEX,
            default_pool_size: 1_000_000,
        });

        let mut mesh_layout = MeshLayoutDesc::new();
        mesh_layout.insert(vertex_attr_id, ());
        mesh_layout.insert(index_attr_id, ());
        let mesh_layout_id = mesh_pool.add_layout(mesh_layout).unwrap();

        let example_vertices = vec![
            Vertex {
                position: [-0.5, 0.5, 0.0],
                color: [1.0, 0.0, 0.0],
            },
            Vertex {
                position: [0.5, 0.5, 0.0],
                color: [0.0, 1.0, 0.0],
            },
            Vertex {
                position: [0.0, -0.5, 0.0],
                color: [0.0, 0.0, 1.0],
            },
        ];

        let example_vertices = AttrBuffer {
            id: vertex_attr_id,
            count: example_vertices.len(),
            data: bytemuck::cast_slice(&example_vertices).to_vec(),
        };

        let example_indices = vec![0u32, 1u32, 2u32];
        let example_indices = AttrBuffer {
            id: index_attr_id,
            count: example_indices.len(),
            data: bytemuck::cast_slice(&example_indices).to_vec(),
        };

        let mut example_mesh = MeshBuffer::default();
        example_mesh.attributes.push(example_vertices);
        example_mesh.attributes.push(example_indices);
        let example_mesh = mesh_pool.load(example_mesh).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 = &[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);

        Self {
            device,
            layouts,
            attr_store,
            shader_info,
            mesh_pool,
            index_attr_id,
            vertex_attr_id,
            mesh_layout_id,
            example_mesh,
            depth_pipeline,
            opaque_pipeline,
            target_info,
        }
    }
}

impl RenderPass for MeshPass {
    type FrameData = FrameData;

    fn create_frame_data(&self) -> FrameData {
        FrameData {}
    }

    fn begin_frame(&self, data: &mut FrameData, phases: &mut Vec<Phase>, queue: &wgpu::Queue) {
        println!("MeshPass::begin_frame()");

        phases.push(Phase::Upload);
        phases.push(Phase::Depth);
        phases.push(Phase::Opaque);
        phases.push(Phase::Transparent);
    }

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

    fn record_render(&self, data: PhaseData<&FrameData>) -> Option<wgpu::RenderBundle> {
        println!("MeshPass::record_render(phase: {:?})", data.phase);

        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: &[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,
                });

        let meshes = &[&self.example_mesh, &self.example_mesh];

        // yikes
        let mesh_bindings: Vec<(MeshLayoutBindings, Vec<(&&MeshHandle, MeshAllocInfos)>)> = self
            .mesh_pool
            .iter_meshes(self.mesh_layout_id, meshes.iter(), |v| v)
            .unwrap()
            .into_iter()
            .map(
                |MeshLayoutInstances {
                     bindings,
                     instances,
                 }| (self.mesh_pool.get_bindings(bindings), instances),
            )
            .collect();

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

        for (bindings, instances) in mesh_bindings.iter() {
            let vertices_pool = bindings.get(self.vertex_attr_id).unwrap();
            let indices_pool = bindings.get(self.index_attr_id).unwrap();

            cmds.set_vertex_buffer(0, vertices_pool.get_buffer().slice(..));
            cmds.set_index_buffer(
                indices_pool.get_buffer().slice(..),
                wgpu::IndexFormat::Uint32,
            );

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

                let is_start = indices.offset as u32;
                let is_end = is_start + indices.count as u32;

                cmds.draw_indexed(is_start..is_end, vertices.offset as i32, 0..1);
                println!(
                    "drew a mesh! {}..{} + {}",
                    is_start, is_end, vertices.offset
                );
            }
        }

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