vidya/crates/renderer/src/images.rs

use std::{borrow::Cow, sync::Arc};

use crate::{buffers::Buffer, define_device_owned_handle, device::DeviceOwned};

use super::{Device, Queue};
use ash::{prelude::*, vk};
use vk_mem::Alloc;

#[derive(Debug, Default, Clone)]
pub struct ImageViewDesc {
    pub flags: vk::ImageViewCreateFlags,
    pub name: Option<Cow<'static, str>>,
    pub kind: vk::ImageViewType,
    pub format: vk::Format,
    pub components: vk::ComponentMapping,
    pub aspect: vk::ImageAspectFlags,
    pub mip_range: MipRange,
    pub layer_range: MipRange,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MipRange(u32, u32);

impl Default for MipRange {
    fn default() -> Self {
        Self(0, vk::REMAINING_ARRAY_LAYERS)
    }
}

impl MipRange {
    fn count(&self) -> u32 {
        self.1
    }
}

impl<R: core::ops::RangeBounds<u32>> From<R> for MipRange {
    fn from(value: R) -> Self {
        let start = match value.start_bound() {
            std::ops::Bound::Included(v) => *v,
            std::ops::Bound::Excluded(v) => *v + 1,
            std::ops::Bound::Unbounded => 0,
        };
        let count = match value.end_bound() {
            std::ops::Bound::Included(v) => *v + 1 - start,
            std::ops::Bound::Excluded(v) => *v - start,
            std::ops::Bound::Unbounded => vk::REMAINING_MIP_LEVELS,
        };

        Self(start, count)
    }
}

impl std::hash::Hash for ImageViewDesc {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.flags.hash(state);
        self.kind.hash(state);
        self.format.hash(state);
        (
            self.components.r,
            self.components.g,
            self.components.b,
            self.components.a,
        )
            .hash(state);
        self.aspect.hash(state);
        self.layer_range.hash(state);
        self.mip_range.hash(state);
    }
}

impl Eq for ImageViewDesc {}
impl PartialEq for ImageViewDesc {
    fn eq(&self, other: &Self) -> bool {
        self.flags == other.flags
            && self.kind == other.kind
            && self.format == other.format
            && (
                self.components.r,
                self.components.g,
                self.components.b,
                self.components.a,
            ) == (
                other.components.r,
                other.components.g,
                other.components.b,
                other.components.a,
            )
            && self.aspect == other.aspect
            && self.mip_range == other.mip_range
            && self.layer_range == other.layer_range
    }
}

#[derive(Debug)]
pub struct Image2D {
    device: Device,
    size: vk::Extent2D,
    mip_levels: u32,
    format: vk::Format,
    image: vk::Image,
    allocation: vk_mem::Allocation,
    name: Option<String>,
}

impl Drop for Image2D {
    fn drop(&mut self) {
        tracing::debug!("destroying image {:?}", self);
        unsafe {
            self.device
                .alloc()
                .destroy_image(self.image, &mut self.allocation);
        }
    }
}

impl Image2D {
    pub fn new_exclusive(
        device: &Device,
        extent: vk::Extent2D,
        mip_levels: u32,
        array_layers: u32,
        format: vk::Format,
        tiling: vk::ImageTiling,
        usage: vk::ImageUsageFlags,
        memory_usage: vk_mem::MemoryUsage,
        alloc_flags: vk_mem::AllocationCreateFlags,
        name: Option<&str>,
    ) -> VkResult<Arc<Self>> {
        let create_info = vk::ImageCreateInfo::default()
            .array_layers(array_layers)
            .mip_levels(mip_levels)
            .extent(vk::Extent3D {
                width: extent.width,
                height: extent.height,
                depth: 1,
            })
            .image_type(vk::ImageType::TYPE_2D)
            .format(format)
            .tiling(tiling)
            .initial_layout(vk::ImageLayout::UNDEFINED)
            .usage(usage)
            .sharing_mode(vk::SharingMode::EXCLUSIVE)
            .samples(vk::SampleCountFlags::TYPE_1);

        let alloc_info = vk_mem::AllocationCreateInfo {
            usage: memory_usage,
            flags: alloc_flags,
            ..Default::default()
        };

        let (image, allocation) =
            unsafe { device.alloc().create_image(&create_info, &alloc_info)? };

        if let Some(name) = name {
            let info = device.alloc().get_allocation_info(&allocation);

            let name = std::ffi::CString::new(name).unwrap_or(c"invalid name".to_owned());
            unsafe {
                device.debug_utils().set_debug_utils_object_name(
                    &vk::DebugUtilsObjectNameInfoEXT::default()
                        .object_handle(info.device_memory)
                        .object_name(&name),
                )?;
            }
        }

        Ok(Arc::new(Self {
            size: extent,
            mip_levels,
            format,
            device: device.clone(),
            image,
            allocation,
            name: name.map(|s| s.to_owned()),
        }))
    }

    pub fn format(&self) -> vk::Format {
        self.format
    }

    pub fn device(&self) -> Device {
        self.device.clone()
    }

    pub fn view(&self, desc: ImageViewDesc) -> VkResult<ImageView> {
        let create_info = vk::ImageViewCreateInfo::default()
            .flags(desc.flags)
            .image(self.image())
            .view_type(vk::ImageViewType::TYPE_2D)
            .format(desc.format)
            .components(desc.components)
            .subresource_range(
                vk::ImageSubresourceRange::default()
                    .aspect_mask(desc.aspect)
                    .base_mip_level(desc.mip_range.0)
                    .level_count(desc.mip_range.count())
                    .base_array_layer(desc.layer_range.0)
                    .layer_count(desc.layer_range.count()),
            );

        let view = unsafe { self.device.dev().create_image_view(&create_info, None)? };

        ImageView::construct(self.device.clone(), view, desc.name)
    }

    pub fn image(&self) -> vk::Image {
        self.image
    }
    pub fn size(&self) -> vk::Extent2D {
        self.size
    }
    pub fn width(&self) -> u32 {
        self.size.width
    }
    pub fn height(&self) -> u32 {
        self.size.height
    }
}

define_device_owned_handle! {
    #[derive(Debug)]
    pub ImageView(vk::ImageView) {} => |this|  unsafe {
        this.device().dev().destroy_image_view(this.handle(), None);
    }
}

pub struct QueueOwnership {
    pub src: u32,
    pub dst: u32,
}

pub fn image_barrier<'a>(
    image: vk::Image,
    aspects: vk::ImageAspectFlags,
    src_stage: vk::PipelineStageFlags2,
    src_access: vk::AccessFlags2,
    dst_stage: vk::PipelineStageFlags2,
    dst_access: vk::AccessFlags2,
    old_layout: vk::ImageLayout,
    new_layout: vk::ImageLayout,
    queue_ownership_op: Option<QueueOwnership>,
) -> vk::ImageMemoryBarrier2<'a> {
    let (src_family, dst_family) = queue_ownership_op
        .map(|t| (t.src, t.dst))
        .unwrap_or((vk::QUEUE_FAMILY_IGNORED, vk::QUEUE_FAMILY_IGNORED));

    vk::ImageMemoryBarrier2::default()
        .image(image)
        .subresource_range(
            vk::ImageSubresourceRange::default()
                .aspect_mask(aspects)
                .base_mip_level(0)
                .base_array_layer(0)
                .level_count(vk::REMAINING_MIP_LEVELS)
                .layer_count(vk::REMAINING_ARRAY_LAYERS),
        )
        .src_stage_mask(src_stage)
        .src_access_mask(src_access)
        .dst_stage_mask(dst_stage)
        .dst_access_mask(dst_access)
        .dst_queue_family_index(dst_family)
        .src_queue_family_index(src_family)
        .old_layout(old_layout)
        .new_layout(new_layout)
}

pub const SUBRESOURCERANGE_COLOR_ALL: vk::ImageSubresourceRange = vk::ImageSubresourceRange {
    aspect_mask: vk::ImageAspectFlags::COLOR,
    base_mip_level: 0,
    level_count: vk::REMAINING_MIP_LEVELS,
    base_array_layer: 0,
    layer_count: vk::REMAINING_ARRAY_LAYERS,
};