vidya/crates/renderer_rendergraph/src/resources.rs

use std::sync::Arc;

use ash::vk;
use renderer_resources::{buffers::Buffer as BufferResource, images::Image as ImageResource};

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

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

impl Buffer {
    pub fn id(self) -> ResourceId {
        self.0
    }
}

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

impl Texture {
    pub fn id(self) -> ResourceId {
        self.0
    }
}

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

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

    pub fn end(self) -> u64 {
        if self.size == vk::WHOLE_SIZE {
            u64::MAX
        } else {
            self.offset.saturating_add(self.size)
        }
    }

    pub fn overlaps(self, other: Self) -> bool {
        self.offset < other.end() && other.offset < self.end()
    }
}

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

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

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

impl SubresourceScope {
    /// Fold identical ranges with compatible access polarity into one scope.
    pub fn try_merge(&mut self, other: &Self) -> bool {
        if access_flag_is_read(self.access) != access_flag_is_read(other.access)
            || access_flag_is_write(self.access) != access_flag_is_write(other.access)
        {
            return false;
        }

        let same_range = match (self.subresource, other.subresource) {
            (SubresourceScopeInner::Buffer(a), SubresourceScopeInner::Buffer(b)) => a == b,
            (
                SubresourceScopeInner::Texture {
                    range: a,
                    layout: layout_a,
                },
                SubresourceScopeInner::Texture {
                    range: b,
                    layout: layout_b,
                },
            ) => texture_ranges_equal(a, b) && layout_a == layout_b,
            _ => false,
        };

        if same_range {
            self.stage |= other.stage;
            self.access |= other.access;
        }
        same_range
    }

    /// Whether accesses overlap and require ordering (a write or layout change).
    pub fn conflicts(&self, other: &Self) -> bool {
        let overlaps = match (self.subresource, other.subresource) {
            (SubresourceScopeInner::Buffer(a), SubresourceScopeInner::Buffer(b)) => a.overlaps(b),
            (
                SubresourceScopeInner::Texture { range: a, .. },
                SubresourceScopeInner::Texture { range: b, .. },
            ) => texture_ranges_overlap(a, b),
            _ => false,
        };

        if !overlaps {
            return false;
        }

        match (self.subresource, other.subresource) {
            (
                SubresourceScopeInner::Texture { layout: a, .. },
                SubresourceScopeInner::Texture { layout: b, .. },
            ) if a != b => true,
            _ => {
                !(access_flag_is_read(self.access)
                    && !access_flag_is_write(self.access)
                    && access_flag_is_read(other.access)
                    && !access_flag_is_write(other.access))
            }
        }
    }
}

fn texture_ranges_equal(a: vk::ImageSubresourceRange, b: vk::ImageSubresourceRange) -> bool {
    a.aspect_mask == b.aspect_mask
        && a.base_mip_level == b.base_mip_level
        && a.level_count == b.level_count
        && a.base_array_layer == b.base_array_layer
        && a.layer_count == b.layer_count
}

fn range_end(base: u32, count: u32) -> u64 {
    if count == vk::REMAINING_MIP_LEVELS || count == vk::REMAINING_ARRAY_LAYERS {
        u64::MAX
    } else {
        u64::from(base) + u64::from(count)
    }
}

fn texture_ranges_overlap(a: vk::ImageSubresourceRange, b: vk::ImageSubresourceRange) -> bool {
    a.aspect_mask.intersects(b.aspect_mask)
        && u64::from(a.base_mip_level) < range_end(b.base_mip_level, b.level_count)
        && u64::from(b.base_mip_level) < range_end(a.base_mip_level, a.level_count)
        && u64::from(a.base_array_layer) < range_end(b.base_array_layer, b.layer_count)
        && u64::from(b.base_array_layer) < range_end(a.base_array_layer, a.layer_count)
}

/// True when an access mask contains a Vulkan read access.
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::ACCELERATION_STRUCTURE_READ_KHR,
    )
}

/// True when an access mask contains a Vulkan write access.
pub fn access_flag_is_write(access: vk::AccessFlags2) -> bool {
    access.intersects(
        vk::AccessFlags2::SHADER_WRITE
            | vk::AccessFlags2::COLOR_ATTACHMENT_WRITE
            | vk::AccessFlags2::DEPTH_STENCIL_ATTACHMENT_WRITE
            | vk::AccessFlags2::TRANSFER_WRITE
            | vk::AccessFlags2::HOST_WRITE
            | vk::AccessFlags2::MEMORY_WRITE
            | vk::AccessFlags2::SHADER_STORAGE_WRITE
            | vk::AccessFlags2::ACCELERATION_STRUCTURE_WRITE_KHR,
    )
}

#[derive(Debug, Clone, Copy)]
pub struct ResourceState {
    pub stage: vk::PipelineStageFlags2,
    pub access: vk::AccessFlags2,
    pub layout: Option<vk::ImageLayout>,
}

impl Default for ResourceState {
    fn default() -> Self {
        Self {
            stage: vk::PipelineStageFlags2::NONE,
            access: vk::AccessFlags2::NONE,
            layout: None,
        }
    }
}

#[derive(Debug, Clone)]
pub enum ResourceDesc {
    CreatedBuffer(renderer_resources::buffers::BufferDesc),
    CreatedTexture(renderer_resources::images::ImageDesc),
    ImportedBuffer {
        resource: Arc<BufferResource>,
        initial: ResourceState,
    },
    ImportedTexture {
        resource: Arc<ImageResource>,
        initial: ResourceState,
    },
}

impl ResourceDesc {
    pub fn initial_state(&self) -> ResourceState {
        match self {
            Self::ImportedBuffer { initial, .. } | Self::ImportedTexture { initial, .. } => {
                *initial
            }
            Self::CreatedBuffer(_) | Self::CreatedTexture(_) => ResourceState::default(),
        }
    }

    pub fn is_created(&self) -> bool {
        matches!(self, Self::CreatedBuffer(_) | Self::CreatedTexture(_))
    }
}