vulkano_test/src/game/assets/square.rs

use std::{error::Error, sync::Arc};

use vulkano::{
    buffer::Subbuffer,
    command_buffer::{AutoCommandBufferBuilder, PrimaryAutoCommandBuffer},
    descriptor_set::{
        DescriptorSet, allocator::StandardDescriptorSetAllocator,
        layout::DescriptorSetLayoutCreateInfo,
    },
    device::Device,
    format::Format,
    memory::allocator::StandardMemoryAllocator,
    pipeline::{
        DynamicState, GraphicsPipeline, Pipeline, PipelineBindPoint, PipelineLayout,
        PipelineShaderStageCreateInfo,
        graphics::{
            GraphicsPipelineCreateInfo,
            color_blend::{ColorBlendAttachmentState, ColorBlendState},
            depth_stencil::{DepthState, DepthStencilState},
            input_assembly::InputAssemblyState,
            multisample::MultisampleState,
            rasterization::RasterizationState,
            subpass::PipelineRenderingCreateInfo,
            vertex_input::{Vertex, VertexDefinition},
            viewport::ViewportState,
        },
        layout::{PipelineDescriptorSetLayoutCreateInfo, PipelineLayoutCreateFlags},
    },
};

use crate::core::render::{
    primitives::{
        AsBindableDescriptorSet, AsDrawable, AsIndexBuffer, AsRecordable, AsVertexBuffer, mvp::Mvp,
        transform::TransformRaw, vertex::Vertex3D,
    },
    texture::Texture,
};

const VERTICES: [Vertex3D; 4] = [
    Vertex3D {
        position: [-0.5, -0.5, 0.0],
        uv: [0.0, 0.0],
    },
    Vertex3D {
        position: [-0.5, 0.5, 0.0],
        uv: [0.0, 1.0],
    },
    Vertex3D {
        position: [0.5, -0.5, 0.0],
        uv: [1.0, 0.0],
    },
    Vertex3D {
        position: [0.5, 0.5, 0.0],
        uv: [1.0, 1.0],
    },
];

const INDICES: [u32; 6] = [0, 2, 1, 2, 3, 1];

pub mod shaders {
    pub mod vs {
        vulkano_shaders::shader! {
            ty: "vertex",
            path: r"res/shaders/vertex.vert",
            generate_structs: false,
        }
    }

    pub mod fs {
        vulkano_shaders::shader! {
            ty: "fragment",
            path: r"res/shaders/vertex.frag",
            generate_structs: false,
        }
    }
}

pub struct Square {
    vertex_buffer: Subbuffer<[Vertex3D]>,
    index_buffer: Subbuffer<[u32]>,
    pipeline: Arc<GraphicsPipeline>,
}

/// Structure pour encapsuler les données de rendu du Square
pub struct SquareRenderData<'a> {
    pub square: &'a Square,
    pub mvp_uniform: &'a Subbuffer<[Mvp]>,
    pub transform_uniform: &'a Subbuffer<[TransformRaw]>,
    pub descriptor_sets: &'a [Arc<DescriptorSet>],
    pub texture: &'a Texture,
}

impl Square {
    pub fn new(
        device: &Arc<Device>,
        memory_allocator: &Arc<StandardMemoryAllocator>,
        swapchain_format: Format,
        depth_format: Format,
    ) -> Result<Self, Box<dyn Error>> {
        let vertex_buffer = Vertex3D::create_vertex_buffer(memory_allocator, &VERTICES)?;

        let index_buffer = u32::create_index_buffer(memory_allocator, &INDICES)?;

        let vs = shaders::vs::load(device.clone())?
            .entry_point("main")
            .ok_or("Failed find main entry point of vertex shader".to_string())?;

        let fs = shaders::fs::load(device.clone())?
            .entry_point("main")
            .ok_or("Failed find main entry point of fragment shader".to_string())?;

        let vertex_input_state =
            [Vertex3D::per_vertex(), TransformRaw::per_instance()].definition(&vs)?;

        let stages = [
            PipelineShaderStageCreateInfo::new(vs),
            PipelineShaderStageCreateInfo::new(fs),
        ];

        let vertex_bindings = Mvp::as_descriptor_set_layout_bindings();
        let texture_bindings = Texture::as_descriptor_set_layout_bindings();

        let vertex_descriptor_set_layout = DescriptorSetLayoutCreateInfo {
            bindings: vertex_bindings,
            ..Default::default()
        };

        let fragment_descriptor_set_layout = DescriptorSetLayoutCreateInfo {
            bindings: texture_bindings,
            ..Default::default()
        };

        let create_info = PipelineDescriptorSetLayoutCreateInfo {
            set_layouts: vec![vertex_descriptor_set_layout, fragment_descriptor_set_layout],
            flags: PipelineLayoutCreateFlags::default(),
            push_constant_ranges: vec![],
        }
        .into_pipeline_layout_create_info(device.clone())?;

        let layout = PipelineLayout::new(device.clone(), create_info)?;

        let subpass = PipelineRenderingCreateInfo {
            color_attachment_formats: vec![Some(swapchain_format)],
            depth_attachment_format: Some(depth_format),
            ..Default::default()
        };

        let pipeline = GraphicsPipeline::new(
            device.clone(),
            None,
            GraphicsPipelineCreateInfo {
                stages: stages.into_iter().collect(),
                vertex_input_state: Some(vertex_input_state),
                input_assembly_state: Some(InputAssemblyState::default()),
                viewport_state: Some(ViewportState::default()),
                rasterization_state: Some(RasterizationState::default()),
                multisample_state: Some(MultisampleState::default()),
                color_blend_state: Some(ColorBlendState::with_attachment_states(
                    subpass.color_attachment_formats.len() as u32,
                    ColorBlendAttachmentState::default(),
                )),
                depth_stencil_state: Some(DepthStencilState {
                    depth: Some(DepthState::simple()),
                    ..Default::default()
                }),
                dynamic_state: [DynamicState::Viewport].into_iter().collect(),
                subpass: Some(subpass.into()),
                ..GraphicsPipelineCreateInfo::layout(layout)
            },
        )?;

        Ok(Self {
            vertex_buffer,
            index_buffer,
            pipeline,
        })
    }

    pub fn render(
        &self,
        command_buffer: &mut AutoCommandBufferBuilder<PrimaryAutoCommandBuffer>,
        descriptor_set_allocator: &Arc<StandardDescriptorSetAllocator>,
        mvp_uniform: &Subbuffer<[Mvp]>,
        transform_uniform: &Subbuffer<[TransformRaw]>,
        texture: &Texture,
    ) -> Result<(), Box<dyn Error>> {
        let layouts = self.pipeline.layout().set_layouts();

        let uniform_descriptor_set =
            Mvp::as_descriptor_set(descriptor_set_allocator, &layouts[0], mvp_uniform)?;

        let texture_descriptor_set =
            Texture::as_descriptor_set(descriptor_set_allocator, &layouts[1], texture)?;

        let render_data = SquareRenderData {
            square: self,
            mvp_uniform,
            transform_uniform,
            texture,
            descriptor_sets: &[uniform_descriptor_set, texture_descriptor_set],
        };

        // Utiliser les nouveaux traits pour le rendu
        Self::record_render_commands(command_buffer, &self.pipeline, &render_data)?;

        Ok(())
    }
}

impl<'a> AsRecordable<SquareRenderData<'a>> for Square {
    fn record_render_commands(
        builder: &mut AutoCommandBufferBuilder<PrimaryAutoCommandBuffer>,
        pipeline: &Arc<GraphicsPipeline>,
        data: &SquareRenderData<'a>,
    ) -> Result<(), Box<dyn Error>> {
        // Bind pipeline
        builder.bind_pipeline_graphics(pipeline.clone())?;

        // Bind descriptor sets
        builder.bind_descriptor_sets(
            PipelineBindPoint::Graphics,
            pipeline.layout().clone(),
            0,
            data.descriptor_sets.iter().cloned().collect::<Vec<_>>(),
        )?;

        // Bind buffers
        builder.bind_vertex_buffers(
            0,
            (
                Self::vertex_buffer(data).clone(),
                data.transform_uniform.clone(),
            ),
        )?;

        if let Some(index_buffer) = Self::index_buffer(data) {
            builder.bind_index_buffer(index_buffer.clone())?;
            unsafe {
                builder.draw_indexed(
                    Self::index_count(data),
                    Self::instance_count(data),
                    0,
                    0,
                    0,
                )?;
            }
        } else {
            unsafe {
                builder.draw(Self::vertex_count(data), Self::instance_count(data), 0, 0)?;
            }
        }

        Ok(())
    }
}