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 side_effect(&self, map: SideEffectMap2);
}

mod impls {
    use super::*;

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

        fn side_effect(&self, mut map: SideEffectMap2) {
            map.insert_range(self.id(), BufferRange::full());
        }
    }

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

        fn side_effect(&self, mut map: SideEffectMap2) {
            map.insert_range(self.id(), self.range);
        }
    }
    impl Resource for BufferSlices {
        fn id(&self) -> ResourceId {
            self.buffer.id()
        }

        fn side_effect(&self, mut map: SideEffectMap2) {
            for range in &self.ranges {
                map.insert_range(self.id(), *range);
            }
        }
    }

    // TODO: the rows are not in bytes, so this is plain wrong. A BufferRowSlice needs to be generic over Buffer/BufferSlice, or commands taking a BufferRowSlice need to take the row parameters separately.
    impl Resource for BufferRowSlice {
        fn id(&self) -> ResourceId {
            self.buffer.id()
        }

        fn side_effect(&self, mut map: SideEffectMap2) {
            map.insert_range(
                self.id(),
                BufferRange {
                    offset: self.row.offset,
                    size: vk::WHOLE_SIZE,
                },
            );
        }
    }
    impl Resource for BufferRowSlices {
        fn id(&self) -> ResourceId {
            self.buffer.id()
        }

        fn side_effect(&self, mut map: SideEffectMap2) {
            for row in &self.rows {
                map.insert_range(
                    self.id(),
                    BufferRange {
                        offset: row.offset,
                        size: vk::WHOLE_SIZE,
                    },
                );
            }
        }
    }

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

        fn side_effect(&self, mut map: SideEffectMap2) {
            map.insert_range(self.id(), TextureRange::full());
        }
    }
    impl Resource for TextureRegion {
        fn id(&self) -> ResourceId {
            self.texture.id()
        }

        fn side_effect(&self, mut map: SideEffectMap2) {
            map.insert_range(self.id(), self.range);
        }
    }
    impl Resource for TextureRegions {
        fn id(&self) -> ResourceId {
            self.texture.id()
        }

        fn side_effect(&self, mut map: SideEffectMap2) {
            for range in &self.ranges {
                map.insert_range(self.id(), *range);
            }
        }
    }
    impl Resource for TextureView {
        fn id(&self) -> ResourceId {
            self.texture.id()
        }

        fn side_effect(&self, mut map: SideEffectMap2) {
            map.insert_range(self.id(), self.range);
        }
    }

    impl<T: Resource> Resource for Read<T> {
        fn id(&self) -> ResourceId {
            self.0.id()
        }

        fn side_effect(&self, map: SideEffectMap2) {
            self.0.side_effect(map);
        }
    }
    impl<T: Resource> Resource for Write<T> {
        fn id(&self) -> ResourceId {
            self.0.id()
        }

        fn side_effect(&self, map: SideEffectMap2) {
            self.0.side_effect(map);
        }
    }
    impl<T: Resource> Resource for ReadWrite<T> {
        fn id(&self) -> ResourceId {
            self.0.id()
        }

        fn side_effect(&self, map: SideEffectMap2) {
            self.0.side_effect(map);
        }
    }
}

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

#[derive(Debug, Default, Clone, Copy)]
pub struct ResourceAccess {
    pub range: ResourceRange,
    pub access: vk::AccessFlags2,
    pub stages: vk::PipelineStageFlags2,
}

impl ResourceAccess {
    /// Attempts to fold another ResourceAccess into this one. Returns true if
    /// successful, false if they are incompatible.
    pub fn try_fold(&mut self, other: &Self) -> bool {
        // if `other` has access flags that aren't in `self`, we can't fold
        // them.
        if !self.access.contains(other.access) {
            return false;
        }

        match (&mut self.range, &other.range) {
            (
                ResourceRange::Buffer { offset, size },
                ResourceRange::Buffer {
                    offset: offset_b,
                    size: size_b,
                },
            ) => {
                // only fold if other lies within self
                if *offset > *offset_b || *offset + *size < *offset_b + *size_b {
                    return false;
                }

                // merge stages: we need the barrier to scope all stages that
                // touch the resource.
                self.stages |= other.stages;

                true
            }
            (
                ResourceRange::Texture {
                    aspect,
                    mip_level,
                    array_layers,
                    layout,
                    ..
                },
                ResourceRange::Texture {
                    aspect: aspect_b,
                    mip_level: mip_level_b,
                    array_layers: array_layers_b,
                    layout: layout_b,
                    ..
                },
            ) => {
                if aspect != aspect_b
                    || layout != layout_b
                    || !mip_level.intersects(mip_level_b)
                    || !array_layers.intersects(array_layers_b)
                {
                    return false;
                }

                *mip_level = mip_level.union(mip_level_b);
                *array_layers = array_layers.union(array_layers_b);
                // let lower_left = Offset::from(*origin).min(Offset::from(*origin_b));
                // let upper_right = (Offset::from(*origin) + Extent::from(*extent).as_offset())
                //     .max(Offset::from(*origin_b) + Extent::from(*extent_b).as_offset());
                // *origin = lower_left.into();
                // *extent = Extent::from_offset(upper_right - lower_left).into();

                // merge stages: we need the barrier to scope all stages that
                // touch the resource.
                self.stages |= other.stages;

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

    pub fn conflicts(&self, other: &Self) -> bool {
        match (&self.range, &other.range) {
            (ResourceRange::Buffer { .. }, ResourceRange::Texture { .. }) => false,
            (ResourceRange::Texture { .. }, ResourceRange::Buffer { .. }) => false,
            (
                ResourceRange::Buffer { offset, size },
                ResourceRange::Buffer {
                    offset: offset_b,
                    size: size_b,
                },
            ) => {
                // two reads don't conflict
                // TODO: two writes may conflict if the user cares about the order of writes.
                if access_flag_is_read(self.access) == access_flag_is_read(other.access) {
                    return false;
                }

                // must be a read/write conflict, check if subresources intersect
                !(*offset > *offset_b + *size_b || *offset_b > *offset + *size)
            }
            (
                ResourceRange::Texture {
                    aspect,
                    mip_level,
                    array_layers,
                    layout,
                    ..
                },
                ResourceRange::Texture {
                    aspect: aspect_b,
                    mip_level: mip_level_b,
                    array_layers: array_layers_b,
                    layout: layout_b,
                    ..
                },
            ) => {
                // layouts differing means we need a layout transition, which is a conflict
                if layout != layout_b {
                    return true;
                }

                // two reads don't conflict
                // TODO: two writes may conflict if the user cares about the order of writes.
                if access_flag_is_read(self.access) == access_flag_is_read(other.access) {
                    return false;
                }

                // must be a read/write conflict, check if subresources intersect
                aspect.intersects(*aspect_b)
                    && (mip_level.intersects(mip_level_b)
                        || array_layers.intersects(array_layers_b))
            }
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub enum ResourceRange {
    Buffer {
        offset: u64,
        size: u64,
    },
    Texture {
        aspect: vk::ImageAspectFlags,
        mip_level: MipRange,
        array_layers: MipRange,
        origin: (i32, i32, i32),
        extent: (u32, u32, u32),
        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,
        }
    }
}

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

#[derive(Debug, Default, Clone, Copy)]
pub struct TextureRange {
    pub aspect: vk::ImageAspectFlags,
    pub mip_levels: MipRange,
    pub array_layers: MipRange,
    pub origin: (i32, i32, i32),
    pub extent: (u32, u32, u32),
}

impl TextureRange {
    pub fn offset(&self) -> vk::Offset3D {
        vk::Offset3D {
            x: self.origin.0,
            y: self.origin.1,
            z: self.origin.2,
        }
    }
    pub fn extent(&self) -> vk::Extent3D {
        vk::Extent3D {
            width: self.extent.0,
            height: self.extent.1,
            depth: self.extent.2,
        }
    }
}

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

impl From<(TextureRange, vk::ImageLayout)> for ResourceRange {
    fn from((value, layout): (TextureRange, vk::ImageLayout)) -> Self {
        ResourceRange::Texture {
            aspect: value.aspect,
            mip_level: value.mip_levels,
            array_layers: value.array_layers,
            origin: value.origin,
            extent: value.extent,
            layout,
        }
    }
}

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

impl TextureRange {
    pub fn full() -> Self {
        Self {
            aspect: vk::ImageAspectFlags::COLOR
                | vk::ImageAspectFlags::DEPTH
                | vk::ImageAspectFlags::STENCIL,
            mip_levels: MipRange::default(),
            array_layers: MipRange::default(),
            origin: (0, 0, 0),
            extent: (u32::MAX, u32::MAX, u32::MAX),
        }
    }
}

pub struct Buffer(pub ResourceId);

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

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 {
    pub texture: Texture,
    pub range: TextureRange,
}

pub struct TextureRegions {
    pub texture: Texture,
    pub ranges: Vec<TextureRange>,
}

pub struct TextureView {
    pub texture: Texture,
    pub format: vk::Format,
    pub mapping: vk::ComponentMapping,
    pub range: TextureRange,
}

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,
    )
}