vidya/crates/renderer/src/render_graph/resources.rs

use std::ops::Deref;

use ash::vk;

use crate::{images::MipRange, render_graph::recorder::SideEffectMap2};

pub trait Resource: Sized {
    fn id(&self) -> ResourceId;
    fn full_range(&self) -> Subresource;
}

mod sealed {
    use ash::vk;

    pub trait Texture {
        fn format(&self) -> vk::Format {
            todo!()
        }
    }
    pub trait Buffer {
        fn size(&self) -> u64 {
            todo!()
        }
    }
}
pub use sealed::{Buffer as AsBuffer, Texture as AsTexture};

#[derive(Debug, Clone, Copy)]
pub enum Subresource {
    Buffer(BufferRange),
    Texture(vk::ImageSubresourceRange),
}

impl Subresource {
    pub fn full_buffer() -> Self {
        Self::Buffer(BufferRange {
            offset: 0,
            size: vk::WHOLE_SIZE,
        })
    }

    pub fn full_texture(format: vk::Format) -> Self {
        use ash_formats::FormatExt;
        Self::Texture(vk::ImageSubresourceRange {
            aspect_mask: format.aspect_flags(),
            base_mip_level: 0,
            level_count: vk::REMAINING_MIP_LEVELS,
            base_array_layer: 0,
            layer_count: vk::REMAINING_ARRAY_LAYERS,
        })
    }
}

#[derive(Debug, Clone, Copy)]
pub struct SubresourceScope {
    pub access: vk::AccessFlags2,
    pub stage: vk::PipelineStageFlags2,
    pub subresource: SubresourceScopeInner,
}

#[derive(Debug, Clone, Copy)]
pub enum SubresourceScopeInner {
    Buffer {
        subresource: BufferRange,
    },
    Texture {
        subresource: vk::ImageSubresourceRange,
        layout: vk::ImageLayout,
    },
}

impl SubresourceScope {
    pub fn try_merge(&mut self, other: &Self) -> bool {
        // cannot merge a read and a write.
        if access_flag_is_read(self.access) != access_flag_is_read(other.access) {
            return false;
        }

        use SubresourceScopeInner::*;
        match (self.subresource, other.subresource) {
            (Buffer { subresource: lhs }, Buffer { subresource: rhs }) => {
                if lhs.offset > rhs.offset || lhs.offset + lhs.size < rhs.offset + rhs.size {
                    // other is not a subresource of self, so we can't merge them.
                    return false;
                }

                self.stage |= other.stage;
                self.access |= other.access;

                true
            }
            (
                Texture {
                    subresource: lhs,
                    layout: lhs_layout,
                },
                Texture {
                    subresource: rhs,
                    layout: rhs_layout,
                },
            ) => {
                // layout transition required even on reads, so can't fold if layouts differ.
                if lhs_layout != rhs_layout {
                    // TODO: consider if we can still merge `other` into `self`
                    // if they overlap, and reduce the second barrier to only a
                    // layout transition.
                    return false;
                }

                // two accesses conflict if they overlap;
                // they overlap if they share mip levels or array layers with
                // the same aspects.
                // If one access contains the other, we can merge them (i
                // think), however: if a previous access requires a less general
                // barrier, then is it better to split?

                let overlaps = lhs.aspect_mask.intersects(rhs.aspect_mask)
                    && lhs.mip_levels().intersects(rhs.mip_levels())
                    && lhs.layers().intersects(rhs.layers());

                // non-overlapping sub-resources shouldn't be merged.
                if !overlaps {
                    return false;
                }

                let lhs_contains_rhs = lhs.aspect_mask.contains(rhs.aspect_mask)
                    && lhs.mip_levels().contains(rhs.mip_levels())
                    && lhs.layers().contains(rhs.layers());

                // other is not a subresource of self, so we can't merge them.
                if !lhs_contains_rhs {
                    return false;
                }

                self.stage |= other.stage;
                self.access |= other.access;

                true
            }
            // different resource types can't merge
            _ => false,
        }
    }

    pub fn conflicts(&self, other: &Self) -> bool {
        use SubresourceScopeInner::*;
        match (self.subresource, other.subresource) {
            (Buffer { subresource: lhs }, Buffer { subresource: rhs }) => {
                // two reads don't conflict
                if access_flag_is_read(self.access) && access_flag_is_read(other.access) {
                    return false;
                }

                // if they overlap, they conflict
                !(lhs.offset > rhs.offset + rhs.size || rhs.offset > lhs.offset + lhs.size)
            }
            (
                Texture {
                    subresource: lhs,
                    layout: lhs_layout,
                },
                Texture {
                    subresource: rhs,
                    layout: rhs_layout,
                },
            ) => {
                // image subresource ranges overlap if they share mip levels and
                // array layers with the same aspects.
                let overlaps = lhs.aspect_mask.intersects(rhs.aspect_mask)
                    && lhs.mip_levels().intersects(rhs.mip_levels())
                    && lhs.layers().intersects(rhs.layers());

                if !overlaps {
                    return false;
                }

                // if they require different layouts, they conflict even on
                // reads, since a layout transition is required.
                if lhs_layout != rhs_layout {
                    return true;
                }

                // finally, two reads don't conflict
                access_flag_is_read(self.access) && access_flag_is_read(other.access)
            }

            // different resource types can't conflict
            _ => false,
        }
    }
}

mod impls {
    use super::*;

    impl sealed::Texture for Texture {}

    impl sealed::Buffer for Buffer {}
    impl sealed::Buffer for BufferSlice {}

    impl Resource for Buffer {
        fn id(&self) -> ResourceId {
            self.0
        }

        fn full_range(&self) -> Subresource {
            Subresource::full_buffer()
        }
    }
    impl Resource for BufferSlice {
        fn id(&self) -> ResourceId {
            self.buffer.id()
        }

        fn full_range(&self) -> Subresource {
            self.buffer.full_range()
        }
    }
    impl Resource for Texture {
        fn id(&self) -> ResourceId {
            self.0
        }
        fn full_range(&self) -> Subresource {
            Subresource::full_texture(vk::Format::UNDEFINED)
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(transparent)]
pub struct ResourceId(pub u64);

#[derive(Debug, Clone, Copy)]
pub enum ResourceRange {
    Buffer {
        offset: u64,
        size: u64,
    },
    Texture {
        aspect: vk::ImageAspectFlags,
        mip_level: MipRange,
        array_layers: MipRange,
        layout: vk::ImageLayout,
    },
}

impl Default for ResourceRange {
    fn default() -> Self {
        ResourceRange::Buffer {
            offset: 0,
            size: u64::MAX,
        }
    }
}

#[derive(Debug, Default, Clone, Copy)]
pub struct BufferRange {
    pub offset: u64,
    pub size: u64,
}

impl BufferRange {
    pub fn full() -> Self {
        Self {
            offset: 0,
            size: vk::WHOLE_SIZE,
        }
    }
}

impl From<BufferRange> for Subresource {
    fn from(value: BufferRange) -> Self {
        Subresource::Buffer(value)
    }
}

#[derive(Debug, Default, Clone, Copy)]
pub struct BufferRows {
    pub offset: u64,
    pub row_count: u32,
    pub row_size: u32,
}

pub trait IntoImageSubresourceLayers {
    fn into_layers(&self) -> vk::ImageSubresourceLayers;
}

pub trait IntoImageSubresourceRange {
    fn into_range(&self) -> vk::ImageSubresourceRange;
}

impl IntoImageSubresourceLayers for vk::ImageSubresource {
    fn into_layers(&self) -> vk::ImageSubresourceLayers {
        vk::ImageSubresourceLayers {
            aspect_mask: self.aspect_mask,
            mip_level: self.mip_level,
            base_array_layer: self.array_layer,
            layer_count: 1,
        }
    }
}

impl<T> IntoImageSubresourceRange for T
where
    T: IntoImageSubresourceLayers,
{
    fn into_range(&self) -> vk::ImageSubresourceRange {
        self.into_layers().into_range()
    }
}

impl IntoImageSubresourceRange for vk::ImageSubresourceLayers {
    fn into_range(&self) -> vk::ImageSubresourceRange {
        vk::ImageSubresourceRange {
            aspect_mask: self.aspect_mask,
            base_mip_level: self.mip_level,
            level_count: 1,
            base_array_layer: self.base_array_layer,
            layer_count: self.layer_count,
        }
    }
}

impl IntoImageSubresourceRange for vk::ImageSubresourceRange {
    fn into_range(&self) -> vk::ImageSubresourceRange {
        *self
    }
}

pub trait ImageSubresourceExt {
    fn aspects(self) -> vk::ImageAspectFlags;
    fn layers(self) -> MipRange;
    fn mip_levels(self) -> MipRange;
}

impl ImageSubresourceExt for vk::ImageSubresourceLayers {
    fn aspects(self) -> vk::ImageAspectFlags {
        self.aspect_mask
    }
    fn layers(self) -> MipRange {
        MipRange::new(
            self.base_array_layer,
            self.base_array_layer + self.layer_count,
        )
    }

    fn mip_levels(self) -> MipRange {
        MipRange::single(self.mip_level)
    }
}

impl ImageSubresourceExt for vk::ImageSubresourceRange {
    fn aspects(self) -> vk::ImageAspectFlags {
        self.aspect_mask
    }
    fn layers(self) -> MipRange {
        MipRange::new(
            self.base_array_layer,
            self.base_array_layer + self.layer_count,
        )
    }

    fn mip_levels(self) -> MipRange {
        MipRange::new(self.base_mip_level, self.base_mip_level + self.level_count)
    }
}

impl ImageSubresourceExt for vk::ImageSubresource {
    fn aspects(self) -> vk::ImageAspectFlags {
        self.aspect_mask
    }
    fn layers(self) -> MipRange {
        MipRange::single(self.mip_level)
    }

    fn mip_levels(self) -> MipRange {
        MipRange::single(self.array_layer)
    }
}

impl From<vk::ImageSubresourceRange> for Subresource {
    fn from(value: vk::ImageSubresourceRange) -> Self {
        Subresource::Texture(value)
    }
}

impl From<(BufferRange, vk::ImageLayout)> for ResourceRange {
    fn from((value, _): (BufferRange, vk::ImageLayout)) -> Self {
        ResourceRange::Buffer {
            offset: value.offset,
            size: value.size,
        }
    }
}

impl From<BufferRange> for ResourceRange {
    fn from(value: BufferRange) -> Self {
        ResourceRange::Buffer {
            offset: value.offset,
            size: value.size,
        }
    }
}

pub struct Buffer(pub ResourceId);

impl Buffer {
    pub fn full_slice(self) -> BufferSlice {
        BufferSlice {
            buffer: self,
            range: BufferRange {
                offset: 0,
                size: u64::MAX,
            },
        }
    }
}

#[derive(Debug, Clone)]
pub struct Texture(pub ResourceId);

pub struct BufferSlice {
    pub buffer: Buffer,
    pub range: BufferRange,
}

pub struct BufferSlices {
    pub buffer: Buffer,
    pub ranges: Vec<BufferRange>,
}

pub struct BufferRowSlice {
    pub buffer: Buffer,
    pub row: BufferRows,
}

pub struct BufferRowSlices {
    pub buffer: Buffer,
    pub rows: Vec<BufferRows>,
}

pub struct TextureRegion<T: IntoImageSubresourceRange = vk::ImageSubresourceRange> {
    pub range: T,
    pub offset: vk::Offset3D,
    pub extent: vk::Extent3D,
}

pub struct Read<T>(pub T);
pub struct Write<T>(pub T);
pub struct ReadWrite<T>(pub T);

impl<T> Deref for Read<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T> Deref for Write<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T> Deref for ReadWrite<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

pub fn access_flag_is_read(access: vk::AccessFlags2) -> bool {
    access.intersects(
        vk::AccessFlags2::INDIRECT_COMMAND_READ
            | vk::AccessFlags2::INDEX_READ
            | vk::AccessFlags2::VERTEX_ATTRIBUTE_READ
            | vk::AccessFlags2::UNIFORM_READ
            | vk::AccessFlags2::INPUT_ATTACHMENT_READ
            | vk::AccessFlags2::SHADER_READ
            | vk::AccessFlags2::COLOR_ATTACHMENT_READ
            | vk::AccessFlags2::DEPTH_STENCIL_ATTACHMENT_READ
            | vk::AccessFlags2::TRANSFER_READ
            | vk::AccessFlags2::HOST_READ
            | vk::AccessFlags2::MEMORY_READ
            | vk::AccessFlags2::SHADER_SAMPLED_READ
            | vk::AccessFlags2::SHADER_STORAGE_READ
            | vk::AccessFlags2::VIDEO_DECODE_READ_KHR
            | vk::AccessFlags2::VIDEO_ENCODE_READ_KHR
            | vk::AccessFlags2::MICROMAP_READ_EXT
            | vk::AccessFlags2::ACCELERATION_STRUCTURE_READ_KHR
            | vk::AccessFlags2::FRAGMENT_DENSITY_MAP_READ_EXT
            | vk::AccessFlags2::FRAGMENT_SHADING_RATE_ATTACHMENT_READ_KHR
            | vk::AccessFlags2::CONDITIONAL_RENDERING_READ_EXT
            | vk::AccessFlags2::SHADER_BINDING_TABLE_READ_KHR
            | vk::AccessFlags2::DESCRIPTOR_BUFFER_READ_EXT
            | vk::AccessFlags2::COLOR_ATTACHMENT_READ_NONCOHERENT_EXT
            | vk::AccessFlags2::COMMAND_PREPROCESS_READ_NV
            | vk::AccessFlags2::TRANSFORM_FEEDBACK_COUNTER_READ_EXT
            | vk::AccessFlags2::INVOCATION_MASK_READ_HUAWEI,
    )
}