vidya/crates/renderer_rendergraph/src/recorder.rs

use std::collections::BTreeMap;

use ash::vk;

use crate::resources::{
    BufferRange, ResourceId, Subresource, SubresourceScope, SubresourceScopeInner,
};

/// The position of a command in its command list.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct CommandIndex(pub usize);

impl CommandIndex {
    pub fn index(self) -> usize {
        self.0
    }
}

impl From<usize> for CommandIndex {
    fn from(index: usize) -> Self {
        Self(index)
    }
}

impl From<CommandIndex> for usize {
    fn from(index: CommandIndex) -> Self {
        index.0
    }
}

#[derive(Debug, Clone, Copy)]
pub struct ResourceAccess {
    pub resource: ResourceId,
    pub scope: SubresourceScope,
}

/// Stores every declared access together with the command that declared it.
#[derive(Debug, Default, Clone)]
pub struct SideEffectMap {
    accesses: Vec<Vec<ResourceAccess>>,
}

impl SideEffectMap {
    pub fn begin_command(&mut self) -> CommandIndex {
        let index = self.accesses.len();
        self.accesses.push(Vec::new());
        index.into()
    }

    pub fn insert(
        &mut self,
        command_index: CommandIndex,
        resource: ResourceId,
        scope: SubresourceScope,
    ) {
        if self.accesses.len() <= command_index.0 {
            self.accesses.resize_with(command_index.0 + 1, Vec::new);
        }
        let accesses = &mut self.accesses[command_index.0];
        for existing in accesses.iter_mut() {
            if existing.resource == resource && existing.scope.try_merge(&scope) {
                return;
            }
        }
        accesses.push(ResourceAccess { resource, scope });
    }

    pub fn accesses(&self, command_index: CommandIndex) -> &[ResourceAccess] {
        self.accesses
            .get(command_index.0)
            .map_or(&[], Vec::as_slice)
    }

    pub fn command_accesses(&self) -> &[Vec<ResourceAccess>] {
        &self.accesses
    }
}

/// A safe, type-erased temporary command buffer. Each command is owned until
/// compilation/recording, while its access declarations are stored separately.
pub struct CommandList {
    pub(crate) commands: Vec<Box<dyn RecordedCommand>>,
    pub side_effects: SideEffectMap,
}

pub trait RecordedCommand {
    fn apply(&mut self, recorder: &mut CommandRecorder<'_>) -> Result<(), RecordError>;
}

impl CommandList {
    pub fn new() -> Self {
        Self {
            commands: Vec::new(),
            side_effects: SideEffectMap::default(),
        }
    }

    pub fn push<C: crate::commands::Command + 'static>(&mut self, command: C) -> CommandIndex {
        let index = self.side_effects.begin_command();
        command.side_effects(index, &mut self.side_effects);
        self.commands.push(Box::new(command));
        index
    }

    pub fn len(&self) -> usize {
        self.commands.len()
    }

    pub fn is_empty(&self) -> bool {
        self.commands.is_empty()
    }
}

impl Default for CommandList {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RecordError {
    MissingBuffer(ResourceId),
    MissingImage(ResourceId),
    UnresolvedCreatedBuffer(ResourceId),
    UnresolvedCreatedImage(ResourceId),
}

/// Vulkan recording context. Handle maps are populated by the compiled graph.
pub struct CommandRecorder<'a> {
    pub device: &'a ash::Device,
    pub command_buffer: vk::CommandBuffer,
    pub buffers: BTreeMap<ResourceId, vk::Buffer>,
    pub images: BTreeMap<ResourceId, vk::Image>,
}

impl<'a> CommandRecorder<'a> {
    pub fn new(device: &'a ash::Device, command_buffer: vk::CommandBuffer) -> Self {
        Self {
            device,
            command_buffer,
            buffers: BTreeMap::new(),
            images: BTreeMap::new(),
        }
    }

    pub fn buffer(&self, id: ResourceId) -> Result<vk::Buffer, RecordError> {
        self.buffers
            .get(&id)
            .copied()
            .ok_or(RecordError::MissingBuffer(id))
    }

    pub fn image(&self, id: ResourceId) -> Result<vk::Image, RecordError> {
        self.images
            .get(&id)
            .copied()
            .ok_or(RecordError::MissingImage(id))
    }

    pub fn pipeline_barrier(&self, barrier: &super::graph::PipelineBarrier) {
        let buffer_barriers: Vec<_> = barrier
            .resources
            .iter()
            .filter_map(|resource| match resource.range {
                Subresource::Buffer(BufferRange { offset, size }) => {
                    self.buffers.get(&resource.id).map(|&buffer| {
                        let (src, dst) = if let Some(queue_transfer) = resource.queue_transfer {
                            (queue_transfer.src.index(), queue_transfer.dst.index())
                        } else {
                            (vk::QUEUE_FAMILY_IGNORED, vk::QUEUE_FAMILY_IGNORED)
                        };

                        vk::BufferMemoryBarrier2::default()
                            .src_stage_mask(resource.src_stage)
                            .src_access_mask(resource.src_access)
                            .dst_stage_mask(resource.dst_stage)
                            .dst_access_mask(resource.dst_access)
                            .src_queue_family_index(src)
                            .dst_queue_family_index(dst)
                            .buffer(buffer)
                            .offset(offset)
                            .size(size)
                    })
                }
                Subresource::Texture(_) => None,
            })
            .collect();

        let image_barriers: Vec<_> = barrier
            .resources
            .iter()
            .filter_map(|resource| match resource.range {
                Subresource::Texture(range) => self.images.get(&resource.id).map(|&image| {
                    vk::ImageMemoryBarrier2::default()
                        .src_stage_mask(resource.src_stage)
                        .src_access_mask(resource.src_access)
                        .dst_stage_mask(resource.dst_stage)
                        .dst_access_mask(resource.dst_access)
                        .old_layout(resource.old_layout.unwrap_or(vk::ImageLayout::UNDEFINED))
                        .new_layout(resource.new_layout.unwrap_or(vk::ImageLayout::GENERAL))
                        .image(image)
                        .subresource_range(range)
                }),
                Subresource::Buffer(_) => None,
            })
            .collect();

        let dependency = vk::DependencyInfo::default()
            .buffer_memory_barriers(&buffer_barriers)
            .image_memory_barriers(&image_barriers);
        unsafe {
            self.device
                .cmd_pipeline_barrier2(self.command_buffer, &dependency);
        }
    }
}

pub(crate) fn scope_range(scope: SubresourceScope) -> Subresource {
    match scope.subresource {
        SubresourceScopeInner::Buffer(range) => Subresource::Buffer(range),
        SubresourceScopeInner::Texture { range, .. } => Subresource::Texture(range),
    }
}