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

//! Опорная сетка и оси — перенос `simv::vk::GridRenderer`.
//!
//! Геометрия строится один раз в буфер, видимый процессору: она статична и крошечная,
//! так что промежуточная копия ради неё не нужна. Дескрипторных наборов нет — всё
//! состояние умещается в одну push-константу.

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

use crate::buffer::{Buffer, MemoryLocation};
use crate::context::Context;
use crate::error::{Ctx, Result};
use crate::shader::{spirv, ShaderModule};

#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct GridVertex {
    pos: Vec3,
    color: Vec3,
}

/// Половина протяжённости сетки в клетках и размер клетки — как в C++-версии.
const HALF_EXTENT_CELLS: i32 = 10;
const CELL_SIZE: f32 = 1.0;

pub struct GridRenderer {
    device: ash::Device,
    layout: vk::PipelineLayout,
    pipeline: vk::Pipeline,
    vertex_buffer: Buffer,
    vertex_count: u32,
    color_format: vk::Format,
    depth_format: vk::Format,
}

impl GridRenderer {
    pub fn new(ctx: &Context) -> Result<Self> {
        let device = ctx.device().clone();

        let ranges = [vk::PushConstantRange::default()
            .stage_flags(vk::ShaderStageFlags::VERTEX)
            .offset(0)
            .size(std::mem::size_of::<Mat4>() as u32)];
        let layout_ci = vk::PipelineLayoutCreateInfo::default().push_constant_ranges(&ranges);
        let layout = unsafe { device.create_pipeline_layout(&layout_ci, None) }
            .ctx("vkCreatePipelineLayout(grid)")?;

        let vertices = build_geometry();
        let bytes: &[u8] = bytemuck::cast_slice(&vertices);
        let mut vertex_buffer = Buffer::new(
            ctx,
            "grid.vertices",
            bytes.len() as vk::DeviceSize,
            vk::BufferUsageFlags::VERTEX_BUFFER,
            MemoryLocation::HostVisible,
        )?;
        vertex_buffer.write(bytes, 0)?;

        Ok(Self {
            device,
            layout,
            pipeline: vk::Pipeline::null(),
            vertex_buffer,
            vertex_count: vertices.len() as u32,
            color_format: vk::Format::UNDEFINED,
            depth_format: vk::Format::UNDEFINED,
        })
    }

    /// Пересобрать конвейер, если форматы вложений сменились (пересоздание swapchain).
    pub fn ensure_pipeline(
        &mut self,
        ctx: &Context,
        color_format: vk::Format,
        depth_format: vk::Format,
    ) -> Result<()> {
        if self.pipeline != vk::Pipeline::null()
            && self.color_format == color_format
            && self.depth_format == depth_format
        {
            return Ok(());
        }
        self.destroy_pipeline();
        self.color_format = color_format;
        self.depth_format = depth_format;
        self.pipeline = self.build_pipeline(ctx)?;
        Ok(())
    }

    fn build_pipeline(&self, ctx: &Context) -> Result<vk::Pipeline> {
        // Модули живут до конца функции и уничтожаются на выходе, включая пути ошибок.
        let vert = ShaderModule::new(&self.device, spirv::GRID_VERT)?;
        let frag = ShaderModule::new(&self.device, spirv::GRID_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"),
        ];

        let bindings = [vk::VertexInputBindingDescription::default()
            .binding(0)
            .stride(std::mem::size_of::<GridVertex>() as u32)
            .input_rate(vk::VertexInputRate::VERTEX)];
        let attributes = [
            vk::VertexInputAttributeDescription::default()
                .location(0)
                .binding(0)
                .format(vk::Format::R32G32B32_SFLOAT)
                .offset(std::mem::offset_of!(GridVertex, pos) as u32),
            vk::VertexInputAttributeDescription::default()
                .location(1)
                .binding(0)
                .format(vk::Format::R32G32B32_SFLOAT)
                .offset(std::mem::offset_of!(GridVertex, color) as u32),
        ];
        let vertex_input = vk::PipelineVertexInputStateCreateInfo::default()
            .vertex_binding_descriptions(&bindings)
            .vertex_attribute_descriptions(&attributes);

        let input_assembly = vk::PipelineInputAssemblyStateCreateInfo::default()
            .topology(vk::PrimitiveTopology::LINE_LIST);
        let viewport = vk::PipelineViewportStateCreateInfo::default()
            .viewport_count(1)
            .scissor_count(1);
        let rasterization = vk::PipelineRasterizationStateCreateInfo::default()
            // Для линейной топологии режим заливки не используется.
            .polygon_mode(vk::PolygonMode::FILL)
            .cull_mode(vk::CullModeFlags::NONE)
            .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::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)| crate::error::VkError::Api {
            context: "vkCreateGraphicsPipelines(grid)",
            source: e,
        })?;
        Ok(pipelines[0])
    }

    pub fn draw(&self, cmd: vk::CommandBuffer, view_proj: &Mat4) {
        if self.pipeline == vk::Pipeline::null() || self.vertex_count == 0 {
            return;
        }
        unsafe {
            self.device
                .cmd_bind_pipeline(cmd, vk::PipelineBindPoint::GRAPHICS, self.pipeline);
            self.device.cmd_push_constants(
                cmd,
                self.layout,
                vk::ShaderStageFlags::VERTEX,
                0,
                bytemuck::bytes_of(view_proj),
            );
            self.device
                .cmd_bind_vertex_buffers(cmd, 0, &[self.vertex_buffer.handle()], &[0]);
            self.device.cmd_draw(cmd, self.vertex_count, 1, 0, 0);
        }
    }

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

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

/// Три плоскости сетки через начало координат плюс цветные оси.
fn build_geometry() -> Vec<GridVertex> {
    const GRID: Vec3 = Vec3::new(0.30, 0.31, 0.35);
    let n = HALF_EXTENT_CELLS;
    let ext = n as f32 * CELL_SIZE;

    let mut verts = Vec::with_capacity(((2 * n + 1) * 6 * 2 + 6) as usize);
    let mut line = |a: Vec3, b: Vec3, color: Vec3| {
        verts.push(GridVertex { pos: a, color });
        verts.push(GridVertex { pos: b, color });
    };

    for i in -n..=n {
        let t = i as f32 * CELL_SIZE;
        // Плоскость XZ (y = 0)
        line(Vec3::new(-ext, 0.0, t), Vec3::new(ext, 0.0, t), GRID);
        line(Vec3::new(t, 0.0, -ext), Vec3::new(t, 0.0, ext), GRID);
        // Плоскость XY (z = 0)
        line(Vec3::new(-ext, t, 0.0), Vec3::new(ext, t, 0.0), GRID);
        line(Vec3::new(t, -ext, 0.0), Vec3::new(t, ext, 0.0), GRID);
        // Плоскость YZ (x = 0)
        line(Vec3::new(0.0, -ext, t), Vec3::new(0.0, ext, t), GRID);
        line(Vec3::new(0.0, t, -ext), Vec3::new(0.0, t, ext), GRID);
    }

    // Оси идут последними, чтобы перекрыть серые центральные линии.
    line(
        Vec3::new(-ext, 0.0, 0.0),
        Vec3::new(ext, 0.0, 0.0),
        Vec3::new(0.85, 0.25, 0.25), // X красная
    );
    line(
        Vec3::new(0.0, -ext, 0.0),
        Vec3::new(0.0, ext, 0.0),
        Vec3::new(0.25, 0.80, 0.30), // Y зелёная
    );
    line(
        Vec3::new(0.0, 0.0, -ext),
        Vec3::new(0.0, 0.0, ext),
        Vec3::new(0.30, 0.45, 0.90), // Z синяя
    );

    verts
}