MineClone/src/render/chunk.rs

use std::collections::HashMap;

use ash::vk;

use crate::{Game, LoadingWorld, RenderCtx, game::chunk::Chunk};

use super::{
    MemType,
    loading_world::WorldComponent,
    mesher::{Mesh, PosNorm},
};

impl WorldComponent for ChunkRender {
    fn descriptors(&self) -> std::collections::HashMap<vk::DescriptorType, u32> {
        let mut map = HashMap::new();
        map.insert(vk::DescriptorType::STORAGE_BUFFER, 1);
        map
    }

    fn desc_layout(&self, ctx: &RenderCtx) -> vk::DescriptorSetLayout {
        setup_desc_layout(ctx)
    }

    fn write_desc_set(
        &self,
        ctx: &RenderCtx,
        desc_set: vk::DescriptorSet,
        skybox: (vk::ImageView, vk::Sampler),
    ) {
        let buf_info = vk::DescriptorBufferInfo::default()
            .buffer(self.mesh_buf)
            .offset(0)
            .range(vk::WHOLE_SIZE);
        let buf_infos = [buf_info];
        let buf_desc = vk::WriteDescriptorSet::default()
            .dst_set(desc_set)
            .dst_binding(0)
            .dst_array_element(0)
            .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
            .buffer_info(&buf_infos);

        unsafe {
            ctx.dev.update_descriptor_sets(&[buf_desc], &[]);
        }
    }

    fn pipeline(&self) -> (vk::Pipeline, vk::PipelineLayout) {
        (self.pipeline, self.pipe_layout)
    }

    fn push_constants(
        &self,
        camera: &crate::game::Camera,
        res: [u32; 2],
        base_color: [f32; 4],
    ) -> Vec<(vk::ShaderStageFlags, u32, Vec<u8>)> {
        let mut constants = Vec::new();
        let cam_data: Vec<u8> = (Vec::<f32>::from([
            camera.origin.x,
            camera.origin.y,
            camera.origin.z,
            0.0,
            camera.rotation.x,
            camera.rotation.y,
            camera.rotation.z,
        ]))
        .iter()
        .map(|f| f.to_ne_bytes())
        .collect::<Vec<_>>()
        .into_flattened();
        constants.push((vk::ShaderStageFlags::VERTEX, 0, cam_data));

        let screen_res: Vec<u8> = res
            .iter()
            .map(|f| f.to_ne_bytes())
            .collect::<Vec<_>>()
            .into_flattened();
        constants.push((vk::ShaderStageFlags::VERTEX, 32, screen_res));

        let base_color: Vec<u8> = (Vec::<f32>::from([0.0, 1.0, 1.0, 1.0]))
            .iter()
            .map(|f| f.to_ne_bytes())
            .collect::<Vec<_>>()
            .into_flattened();
        constants.push((vk::ShaderStageFlags::FRAGMENT, 48, base_color));

        constants
    }

    fn num_verticies(&self) -> u32 {
        self.num_vert as u32
    }
}

pub struct ChunkRender {
    mesh_buf: vk::Buffer,
    num_vert: usize,
    pipeline: vk::Pipeline,
    pipe_layout: vk::PipelineLayout,
}

impl RenderCtx {
    fn create_chunk(&self, pass: vk::RenderPass, chunk: &Chunk) -> ChunkRender {
        let mut mesh = Mesh::default();
        chunk.mesh(&mut mesh);
        let num_vert = mesh.num_vert();
        let buf = mesh.to_buffer();

        let mesh_mem = self.mem_alloc(
            MemType::HostVisibile,
            (buf.len() * size_of::<PosNorm>()) as u64,
        );

        let mesh_buf = self.alloc_buf(
            mesh_mem,
            (buf.len() * size_of::<PosNorm>()) as u64,
            vk::BufferUsageFlags::TRANSFER_SRC | vk::BufferUsageFlags::STORAGE_BUFFER,
        );

        let mesh_mem = self.buf_to_ptr(mesh_mem, mesh_buf);
        unsafe {
            std::ptr::copy(buf.as_ptr(), mesh_mem as *mut PosNorm, buf.len());
        }

        let pipe_layout = setup_pipe_layout(self);
        let pipeline = setup_pipeline(self, &pass, pipe_layout);

        ChunkRender {
            mesh_buf,
            num_vert,
            pipeline,
            pipe_layout,
        }
    }
}

impl Game<LoadingWorld> {
    /// Create a cube to be rendered in world, this can then be
    /// added using [`Game<LoadingWorld>::add_component`]
    pub fn create_chunk(&self, chunk: &Chunk) -> ChunkRender {
        let Self {
            state: LoadingWorld { ctx, pass, .. },
        } = self;
        ctx.create_chunk(*pass, chunk)
    }
}

fn load_chunk_shaders(ctx: &RenderCtx) -> (vk::ShaderModule, vk::ShaderModule) {
    let shader_vert_shader = ctx.shader_mod_from_file(
        "shaders/chunk_vert.spv",
        vk::ShaderModuleCreateFlags::empty(),
    );
    let shader_frag_shader = ctx.shader_mod_from_file(
        "shaders/chunk_frag.spv",
        vk::ShaderModuleCreateFlags::empty(),
    );

    (shader_vert_shader, shader_frag_shader)
}

fn setup_desc_layout(ctx: &RenderCtx) -> vk::DescriptorSetLayout {
    let storage_binding = vk::DescriptorSetLayoutBinding::default()
        .binding(0)
        .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
        .descriptor_count(1)
        .stage_flags(vk::ShaderStageFlags::VERTEX);
    let layouts = [storage_binding];

    let layout_info = vk::DescriptorSetLayoutCreateInfo::default().bindings(&layouts);
    let layout = unsafe {
        ctx.dev
            .create_descriptor_set_layout(&layout_info, None)
            .expect("Failed to create descriptor set layout")
    };
    layout
}

fn setup_pipe_layout(ctx: &RenderCtx) -> vk::PipelineLayout {
    let layout = setup_desc_layout(ctx);
    let push_range: [vk::PushConstantRange; 2] = [
        vk::PushConstantRange::default()
            .stage_flags(vk::ShaderStageFlags::VERTEX)
            .offset(0)
            .size(48), // vec4 camera_orig, camera_rot; uvec2 screen_res
        vk::PushConstantRange::default()
            .stage_flags(vk::ShaderStageFlags::FRAGMENT)
            .offset(48)
            .size(16), // vec4 base_color
    ];

    let layouts = &[layout];

    let pipe_layout = vk::PipelineLayoutCreateInfo::default()
        .set_layouts(layouts.as_ref())
        .push_constant_ranges(&push_range);
    let pipe_layout = unsafe {
        ctx.dev
            .create_pipeline_layout(&pipe_layout, None)
            .expect("Failed to create pipeline layout")
    };
    pipe_layout
}

fn setup_pipeline(
    ctx: &RenderCtx,
    pass: &vk::RenderPass,
    layout: vk::PipelineLayout,
) -> vk::Pipeline {
    let (vert_shader, frag_shader) = load_chunk_shaders(ctx);
    let shader_stages = ctx.setup_simple_shader_stage(vert_shader, frag_shader);
    let vert_input = vk::PipelineVertexInputStateCreateInfo::default();
    let input_assembly = vk::PipelineInputAssemblyStateCreateInfo::default()
        .topology(vk::PrimitiveTopology::TRIANGLE_LIST);
    let rasterization = vk::PipelineRasterizationStateCreateInfo::default()
        .polygon_mode(vk::PolygonMode::FILL)
        .cull_mode(vk::CullModeFlags::BACK)
        .front_face(vk::FrontFace::CLOCKWISE)
        .line_width(1.0);
    let multisample = vk::PipelineMultisampleStateCreateInfo::default()
        .rasterization_samples(vk::SampleCountFlags::TYPE_1);
    let depth_stencil = vk::PipelineDepthStencilStateCreateInfo::default()
        .depth_test_enable(true)
        .depth_write_enable(true)
        .depth_compare_op(vk::CompareOp::GREATER)
        .depth_bounds_test_enable(false)
        .stencil_test_enable(false)
        .min_depth_bounds(0.0)
        .max_depth_bounds(1.0);

    let blend = vk::PipelineColorBlendAttachmentState::default()
        .src_color_blend_factor(vk::BlendFactor::SRC_COLOR)
        .dst_color_blend_factor(vk::BlendFactor::SRC_COLOR)
        .color_blend_op(vk::BlendOp::ADD)
        .src_alpha_blend_factor(vk::BlendFactor::SRC_COLOR)
        .dst_alpha_blend_factor(vk::BlendFactor::SRC_COLOR)
        .alpha_blend_op(vk::BlendOp::ADD)
        .color_write_mask(
            vk::ColorComponentFlags::R
                | vk::ColorComponentFlags::G
                | vk::ColorComponentFlags::B
                | vk::ColorComponentFlags::A,
        );

    let blend_attachments = [blend];

    let color_blend = vk::PipelineColorBlendStateCreateInfo::default()
        .attachments(blend_attachments.as_ref())
        .logic_op_enable(false)
        .logic_op(vk::LogicOp::COPY)
        .blend_constants([1.0, 1.0, 1.0, 1.0]);

    let dynamic = vk::PipelineDynamicStateCreateInfo::default().dynamic_states(&[
        vk::DynamicState::VIEWPORT_WITH_COUNT,
        vk::DynamicState::SCISSOR_WITH_COUNT,
    ]);

    let pipe_info = vk::GraphicsPipelineCreateInfo::default()
        .stages(&shader_stages)
        .vertex_input_state(&vert_input)
        .input_assembly_state(&input_assembly)
        .rasterization_state(&rasterization)
        .multisample_state(&multisample)
        .depth_stencil_state(&depth_stencil)
        .color_blend_state(&color_blend)
        .layout(layout)
        .dynamic_state(&dynamic)
        .render_pass(*pass)
        .subpass(0);

    let pipe = unsafe {
        ctx.dev
            .create_graphics_pipelines(vk::PipelineCache::null(), &[pipe_info], None)
            .expect("Failed to create graphics pipeline")
            .remove(0)
    };
    println!("Created pipeline");
    pipe
}