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

//! Изображение с представлением — перенос `simv::vk::Image`.
//!
//! Как и `buffer.rs`, в C++-версии лежит без применения: глубинное вложение
//! `Renderer` создаёт там свои `VkImage`/`VkImageView` напрямую. Здесь на нём как раз
//! и построено глубинное вложение — см. заметку в `buffer.rs`.

use std::sync::{Arc, Mutex};

use ash::vk;
use gpu_allocator::vulkan::{Allocation, AllocationCreateDesc, AllocationScheme, Allocator};

use crate::context::Context;
use crate::error::{Ctx, Result, VkError};

pub struct ImageDesc {
    pub extent: vk::Extent3D,
    pub format: vk::Format,
    pub usage: vk::ImageUsageFlags,
    pub image_type: vk::ImageType,
    pub view_type: vk::ImageViewType,
    pub aspect: vk::ImageAspectFlags,
}

impl Default for ImageDesc {
    fn default() -> Self {
        Self {
            extent: vk::Extent3D {
                width: 1,
                height: 1,
                depth: 1,
            },
            format: vk::Format::R8G8B8A8_UNORM,
            usage: vk::ImageUsageFlags::SAMPLED,
            image_type: vk::ImageType::TYPE_2D,
            view_type: vk::ImageViewType::TYPE_2D,
            aspect: vk::ImageAspectFlags::COLOR,
        }
    }
}

pub struct Image {
    device: ash::Device,
    allocator: Arc<Mutex<Allocator>>,
    image: vk::Image,
    view: vk::ImageView,
    allocation: Option<Allocation>,
    extent: vk::Extent3D,
    format: vk::Format,
    aspect: vk::ImageAspectFlags,
    /// Текущая раскладка. Отслеживается здесь, чтобы `transition_to` не выставлял
    /// барьер на переход в ту же раскладку — так же, как в C++-версии.
    layout: vk::ImageLayout,
}

impl Image {
    pub fn new(ctx: &Context, name: &str, desc: &ImageDesc) -> Result<Self> {
        let device = ctx.device().clone();
        let ci = vk::ImageCreateInfo::default()
            .image_type(desc.image_type)
            .format(desc.format)
            .extent(desc.extent)
            .mip_levels(1)
            .array_layers(1)
            .samples(vk::SampleCountFlags::TYPE_1)
            .tiling(vk::ImageTiling::OPTIMAL)
            .usage(desc.usage)
            .sharing_mode(vk::SharingMode::EXCLUSIVE)
            .initial_layout(vk::ImageLayout::UNDEFINED);
        let image = unsafe { device.create_image(&ci, None) }.ctx("vkCreateImage")?;
        let requirements = unsafe { device.get_image_memory_requirements(image) };

        let allocator = ctx.allocator().clone();
        let allocation = {
            let mut guard = allocator.lock().expect("аллокатор отравлен паникой");
            guard.allocate(&AllocationCreateDesc {
                name,
                requirements,
                location: gpu_allocator::MemoryLocation::GpuOnly,
                linear: false,
                allocation_scheme: AllocationScheme::GpuAllocatorManaged,
            })
        };
        let allocation = match allocation {
            Ok(a) => a,
            Err(e) => {
                unsafe { device.destroy_image(image, None) };
                return Err(VkError::Other(format!("gpu-allocator: {e}")));
            }
        };
        unsafe { device.bind_image_memory(image, allocation.memory(), allocation.offset()) }
            .ctx("vkBindImageMemory")?;

        let view_ci = vk::ImageViewCreateInfo::default()
            .image(image)
            .view_type(desc.view_type)
            .format(desc.format)
            .subresource_range(vk::ImageSubresourceRange {
                aspect_mask: desc.aspect,
                base_mip_level: 0,
                level_count: 1,
                base_array_layer: 0,
                layer_count: 1,
            });
        let view = unsafe { device.create_image_view(&view_ci, None) }.ctx("vkCreateImageView")?;

        Ok(Self {
            device,
            allocator,
            image,
            view,
            allocation: Some(allocation),
            extent: desc.extent,
            format: desc.format,
            aspect: desc.aspect,
            layout: vk::ImageLayout::UNDEFINED,
        })
    }

    pub fn image(&self) -> vk::Image {
        self.image
    }
    pub fn view(&self) -> vk::ImageView {
        self.view
    }
    pub fn extent(&self) -> vk::Extent3D {
        self.extent
    }
    pub fn format(&self) -> vk::Format {
        self.format
    }
    pub fn layout(&self) -> vk::ImageLayout {
        self.layout
    }

    /// Барьер перехода раскладки на sync2. Повторный переход в ту же раскладку —
    /// не операция.
    pub fn transition_to(
        &mut self,
        cmd: vk::CommandBuffer,
        new_layout: vk::ImageLayout,
        src_stage: vk::PipelineStageFlags2,
        src_access: vk::AccessFlags2,
        dst_stage: vk::PipelineStageFlags2,
        dst_access: vk::AccessFlags2,
    ) {
        if new_layout == self.layout {
            return;
        }
        let barrier = vk::ImageMemoryBarrier2::default()
            .src_stage_mask(src_stage)
            .src_access_mask(src_access)
            .dst_stage_mask(dst_stage)
            .dst_access_mask(dst_access)
            .old_layout(self.layout)
            .new_layout(new_layout)
            .image(self.image)
            .subresource_range(vk::ImageSubresourceRange {
                aspect_mask: self.aspect,
                base_mip_level: 0,
                level_count: 1,
                base_array_layer: 0,
                layer_count: 1,
            });
        let barriers = [barrier];
        let dependency = vk::DependencyInfo::default().image_memory_barriers(&barriers);
        unsafe { self.device.cmd_pipeline_barrier2(cmd, &dependency) };
        self.layout = new_layout;
    }
}

impl Drop for Image {
    fn drop(&mut self) {
        unsafe { self.device.destroy_image_view(self.view, None) };
        if let Some(allocation) = self.allocation.take() {
            let _ = self
                .allocator
                .lock()
                .map(|mut guard| guard.free(allocation));
        }
        unsafe { self.device.destroy_image(self.image, None) };
    }
}