CFDManager/rust/crates/simv-vk/src/mesh_renderer.rs

//! Конвейеры модели — перенос `simv::vk::MeshRenderer`.
//!
//! Два конвейера на один макет: заливка и `POLYGON_MODE_LINE` для каркаса. У линейного
//! включено небольшое смещение глубины, чтобы каркас лёг поверх заливки, а не спорил
//! с ней за z. Отсечение выключено — у загруженных моделей обход граней вперемешку.
//!
//! ВАЖНО: `MeshPushConstants` обязан байт в байт совпадать с блоком `push_constant`
//! в `shaders/editor/mesh.vert` и `mesh.frag`. Поле `color.a` — это ФЛАГ (1 — плоская
//! заливка, 0 — постоянный цвет каркаса), а не прозрачность. Меняете здесь — меняйте
//! и в шейдере.

use ash::vk;
use glam::{Mat4, Vec4};

use crate::context::Context;
use crate::error::{Ctx, Result, VkError};
use crate::gpu_mesh::GpuMesh;
use crate::render_mode::RenderMode;
use crate::shader::{spirv, ShaderModule};

#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct MeshPushConstants {
    mvp: Mat4,
    color: Vec4,
}

/// `color.a`: 1 — плоская заливка, 0 — постоянный цвет.
const SOLID: Vec4 = Vec4::new(0.72, 0.74, 0.78, 1.0);
const WIRE: Vec4 = Vec4::new(0.78, 0.82, 0.88, 0.0);
const OVERLAY: Vec4 = Vec4::new(0.06, 0.07, 0.09, 0.0);

pub struct MeshRenderer {
    device: ash::Device,
    layout: vk::PipelineLayout,
    fill: vk::Pipeline,
    line: vk::Pipeline,
    color_format: vk::Format,
    depth_format: vk::Format,
}

impl MeshRenderer {
    pub fn new(ctx: &Context) -> Result<Self> {
        let device = ctx.device().clone();
        let ranges = [vk::PushConstantRange::default()
            .stage_flags(vk::ShaderStageFlags::VERTEX | vk::ShaderStageFlags::FRAGMENT)
            .offset(0)
            .size(std::mem::size_of::<MeshPushConstants>() as u32)];
        let layout_ci = vk::PipelineLayoutCreateInfo::default().push_constant_ranges(&ranges);
        let layout = unsafe { device.create_pipeline_layout(&layout_ci, None) }
            .ctx("vkCreatePipelineLayout(mesh)")?;

        Ok(Self {
            device,
            layout,
            fill: vk::Pipeline::null(),
            line: vk::Pipeline::null(),
            color_format: vk::Format::UNDEFINED,
            depth_format: vk::Format::UNDEFINED,
        })
    }

    pub fn ensure_pipelines(
        &mut self,
        ctx: &Context,
        color_format: vk::Format,
        depth_format: vk::Format,
    ) -> Result<()> {
        if self.fill != vk::Pipeline::null()
            && self.color_format == color_format
            && self.depth_format == depth_format
        {
            return Ok(());
        }
        self.destroy_pipelines();
        self.color_format = color_format;
        self.depth_format = depth_format;
        self.fill = self.build_pipeline(ctx, vk::PolygonMode::FILL, false)?;
        self.line = self.build_pipeline(ctx, vk::PolygonMode::LINE, true)?;
        Ok(())
    }

    fn build_pipeline(
        &self,
        ctx: &Context,
        polygon_mode: vk::PolygonMode,
        depth_bias: bool,
    ) -> Result<vk::Pipeline> {
        let vert = ShaderModule::new(&self.device, spirv::MESH_VERT)?;
        let frag = ShaderModule::new(&self.device, spirv::MESH_FRAG)?;
        let stages = [
            vk::PipelineShaderStageCreateInfo::default()
                .stage(vk::ShaderStageFlags::VERTEX)
                .module(vert.handle())
                .name(c"main"),
            vk::PipelineShaderStageCreateInfo::default()
                .stage(vk::ShaderStageFlags::FRAGMENT)
                .module(frag.handle())
                .name(c"main"),
        ];

        // Поток вершин — плотный float3: `mesh.frag` восстанавливает нормаль грани из
        // экранных производных, так что нормали передавать не нужно.
        let bindings = [vk::VertexInputBindingDescription::default()
            .binding(0)
            .stride(std::mem::size_of::<f32>() as u32 * 3)
            .input_rate(vk::VertexInputRate::VERTEX)];
        let attributes = [vk::VertexInputAttributeDescription::default()
            .location(0)
            .binding(0)
            .format(vk::Format::R32G32B32_SFLOAT)
            .offset(0)];
        let vertex_input = vk::PipelineVertexInputStateCreateInfo::default()
            .vertex_binding_descriptions(&bindings)
            .vertex_attribute_descriptions(&attributes);

        let input_assembly = vk::PipelineInputAssemblyStateCreateInfo::default()
            .topology(vk::PrimitiveTopology::TRIANGLE_LIST);
        let viewport = vk::PipelineViewportStateCreateInfo::default()
            .viewport_count(1)
            .scissor_count(1);

        let mut rasterization = vk::PipelineRasterizationStateCreateInfo::default()
            .polygon_mode(polygon_mode)
            // У загруженных моделей обход граней бывает вперемешку.
            .cull_mode(vk::CullModeFlags::NONE)
            .front_face(vk::FrontFace::COUNTER_CLOCKWISE)
            .line_width(1.0);
        if depth_bias {
            // Подвинуть каркас к камере, чтобы он лёг поверх заливки.
            rasterization = rasterization
                .depth_bias_enable(true)
                .depth_bias_constant_factor(-1.0)
                .depth_bias_slope_factor(-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::LESS_OR_EQUAL)
            .min_depth_bounds(0.0)
            .max_depth_bounds(1.0);
        let blend_attachments = [vk::PipelineColorBlendAttachmentState::default()
            .color_write_mask(vk::ColorComponentFlags::RGBA)];
        let color_blend =
            vk::PipelineColorBlendStateCreateInfo::default().attachments(&blend_attachments);
        let dynamic_states = [vk::DynamicState::VIEWPORT, vk::DynamicState::SCISSOR];
        let dynamic = vk::PipelineDynamicStateCreateInfo::default().dynamic_states(&dynamic_states);

        let color_formats = [self.color_format];
        let mut rendering = vk::PipelineRenderingCreateInfo::default()
            .color_attachment_formats(&color_formats)
            .depth_attachment_format(self.depth_format);

        let create_info = [vk::GraphicsPipelineCreateInfo::default()
            .stages(&stages)
            .vertex_input_state(&vertex_input)
            .input_assembly_state(&input_assembly)
            .viewport_state(&viewport)
            .rasterization_state(&rasterization)
            .multisample_state(&multisample)
            .depth_stencil_state(&depth_stencil)
            .color_blend_state(&color_blend)
            .dynamic_state(&dynamic)
            .layout(self.layout)
            .push_next(&mut rendering)];

        let pipelines = unsafe {
            self.device
                .create_graphics_pipelines(ctx.pipeline_cache(), &create_info, None)
        }
        .map_err(|(_, e)| VkError::Api {
            context: "vkCreateGraphicsPipelines(mesh)",
            source: e,
        })?;
        Ok(pipelines[0])
    }

    pub fn draw(&self, cmd: vk::CommandBuffer, mesh: &GpuMesh, mvp: &Mat4, mode: RenderMode) {
        if self.fill == vk::Pipeline::null() || self.line == vk::Pipeline::null() {
            return;
        }
        unsafe {
            self.device
                .cmd_bind_vertex_buffers(cmd, 0, &[mesh.vertex_buffer()], &[0]);
            self.device
                .cmd_bind_index_buffer(cmd, mesh.index_buffer(), 0, vk::IndexType::UINT32);
        }

        let draw_with = |pipeline: vk::Pipeline, color: Vec4| {
            let push = MeshPushConstants { mvp: *mvp, color };
            unsafe {
                self.device
                    .cmd_bind_pipeline(cmd, vk::PipelineBindPoint::GRAPHICS, pipeline);
                self.device.cmd_push_constants(
                    cmd,
                    self.layout,
                    vk::ShaderStageFlags::VERTEX | vk::ShaderStageFlags::FRAGMENT,
                    0,
                    bytemuck::bytes_of(&push),
                );
                self.device
                    .cmd_draw_indexed(cmd, mesh.index_count(), 1, 0, 0, 0);
            }
        };

        match mode {
            RenderMode::Solid => draw_with(self.fill, SOLID),
            RenderMode::Wireframe => draw_with(self.line, WIRE),
            RenderMode::SolidWire => {
                draw_with(self.fill, SOLID);
                draw_with(self.line, OVERLAY);
            }
        }
    }

    fn destroy_pipelines(&mut self) {
        unsafe {
            if self.fill != vk::Pipeline::null() {
                self.device.destroy_pipeline(self.fill, None);
                self.fill = vk::Pipeline::null();
            }
            if self.line != vk::Pipeline::null() {
                self.device.destroy_pipeline(self.line, None);
                self.line = vk::Pipeline::null();
            }
        }
    }
}

impl Drop for MeshRenderer {
    fn drop(&mut self) {
        self.destroy_pipelines();
        unsafe { self.device.destroy_pipeline_layout(self.layout, None) };
    }
}