vidya/crates/renderer_rendergraph/src/graph.rs

use std::collections::{BTreeMap, BTreeSet};

use ash::vk;

use crate::{
    recorder::{CommandIndex, CommandList, CommandRecorder, RecordError, scope_range},
    resources::{
        Buffer, ResourceDesc, ResourceId, Subresource, SubresourceScope, SubresourceScopeInner,
        Texture, access_flag_is_write,
    },
};

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GraphError {
    UnknownResource(ResourceId),
    DependencyCycle,
}

#[derive(Debug, Clone, Copy)]
pub struct BarrierResource {
    pub id: ResourceId,
    pub range: Subresource,
    pub src_stage: vk::PipelineStageFlags2,
    pub src_access: vk::AccessFlags2,
    pub dst_stage: vk::PipelineStageFlags2,
    pub dst_access: vk::AccessFlags2,
    pub old_layout: Option<vk::ImageLayout>,
    pub new_layout: Option<vk::ImageLayout>,
}

#[derive(Debug, Clone)]
pub struct PipelineBarrier {
    pub before_command: CommandIndex,
    pub resources: Vec<BarrierResource>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ResourceLifetime {
    pub first_command: CommandIndex,
    pub last_command: CommandIndex,
    /// Placeholder allocator binding; resources with disjoint lifetimes can
    /// share an offset when their descriptor sizes permit it.
    pub memory_offset: u64,
}

pub struct RenderGraph {
    next_id: u64,
    resources: BTreeMap<ResourceId, ResourceDesc>,
    pub commands: CommandList,
}

impl RenderGraph {
    pub fn new() -> Self {
        Self {
            next_id: 0,
            resources: BTreeMap::new(),
            commands: CommandList::new(),
        }
    }

    fn add_resource(&mut self, desc: ResourceDesc) -> ResourceId {
        let id = ResourceId(self.next_id);
        self.next_id = self.next_id.checked_add(1).expect("resource id overflow");
        self.resources.insert(id, desc);
        id
    }

    pub fn create_buffer(&mut self, desc: crate::resources::BufferDesc) -> Buffer {
        Buffer(self.add_resource(ResourceDesc::CreatedBuffer(desc)))
    }

    pub fn create_texture(&mut self, desc: crate::resources::TextureDesc) -> Texture {
        Texture(self.add_resource(ResourceDesc::CreatedTexture(desc)))
    }

    pub fn import_buffer(
        &mut self,
        handle: vk::Buffer,
        initial: crate::resources::ResourceState,
    ) -> Buffer {
        Buffer(self.add_resource(ResourceDesc::ImportedBuffer { handle, initial }))
    }

    pub fn import_texture(
        &mut self,
        handle: vk::Image,
        initial: crate::resources::ResourceState,
    ) -> Texture {
        Texture(self.add_resource(ResourceDesc::ImportedTexture { handle, initial }))
    }

    pub fn push<C: crate::commands::Command + 'static>(&mut self, command: C) -> CommandIndex {
        self.commands.push(command)
    }

    pub fn compile(self) -> Result<CompiledGraph, GraphError> {
        let command_count = self.commands.len();
        let mut edges = vec![BTreeSet::new(); command_count];
        let mut indegree = vec![0usize; command_count];
        let accesses = self.commands.side_effects.command_accesses();

        for current in 0..command_count {
            for prior in 0..current {
                let dependency = accesses[prior].iter().any(|earlier| {
                    accesses[current].iter().any(|later| {
                        earlier.resource == later.resource && earlier.scope.conflicts(&later.scope)
                    })
                });
                if dependency && edges[prior].insert(current) {
                    indegree[current] += 1;
                }
            }
        }

        let mut ready: BTreeSet<_> = indegree
            .iter()
            .enumerate()
            .filter_map(|(index, &degree)| (degree == 0).then_some(index))
            .collect();
        let mut order = Vec::with_capacity(command_count);
        while let Some(index) = ready.pop_first() {
            order.push(index.into());
            for &dependent in &edges[index] {
                indegree[dependent] -= 1;
                if indegree[dependent] == 0 {
                    ready.insert(dependent);
                }
            }
        }
        if order.len() != command_count {
            return Err(GraphError::DependencyCycle);
        }

        for command_accesses in accesses {
            for access in command_accesses {
                if !self.resources.contains_key(&access.resource) {
                    return Err(GraphError::UnknownResource(access.resource));
                }
            }
        }

        let barriers = build_barriers(&order, accesses, &self.resources);
        let lifetimes = bind_resources(&self.resources, accesses);

        Ok(CompiledGraph {
            resources: self.resources,
            commands: self.commands,
            order,
            barriers,
            lifetimes,
        })
    }
}

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

pub struct CompiledGraph {
    resources: BTreeMap<ResourceId, ResourceDesc>,
    commands: CommandList,
    order: Vec<CommandIndex>,
    barriers: Vec<PipelineBarrier>,
    lifetimes: BTreeMap<ResourceId, ResourceLifetime>,
}

impl CompiledGraph {
    pub fn command_order(&self) -> &[CommandIndex] {
        &self.order
    }

    pub fn barriers(&self) -> &[PipelineBarrier] {
        &self.barriers
    }

    pub fn resource_lifetimes(&self) -> &BTreeMap<ResourceId, ResourceLifetime> {
        &self.lifetimes
    }

    /// Record the compiled order into a real Vulkan command buffer. Imported
    /// handles are automatically bound. Created resources require a future
    /// allocator integration and currently return a missing-handle error when
    /// used by a command.
    pub fn record_commands(
        &mut self,
        recorder: &mut CommandRecorder<'_>,
    ) -> Result<(), RecordError> {
        for (&id, &desc) in &self.resources {
            match desc {
                ResourceDesc::ImportedBuffer { handle, .. } => {
                    recorder.buffers.insert(id, handle);
                }
                ResourceDesc::ImportedTexture { handle, .. } => {
                    recorder.images.insert(id, handle);
                }
                ResourceDesc::CreatedBuffer(_) | ResourceDesc::CreatedTexture(_) => {}
            }
        }

        let mut barriers = self.barriers.iter().peekable();
        for (ordered_index, &command_index) in self.order.iter().enumerate() {
            while barriers
                .peek()
                .is_some_and(|barrier| barrier.before_command == ordered_index.into())
            {
                recorder.pipeline_barrier(barriers.next().expect("peeked barrier"));
            }
            self.commands.commands[command_index.index()].apply(recorder)?;
        }
        Ok(())
    }
}

fn build_barriers(
    order: &[CommandIndex],
    accesses: &[Vec<crate::recorder::ResourceAccess>],
    resource_descs: &BTreeMap<ResourceId, ResourceDesc>,
) -> Vec<PipelineBarrier> {
    let mut result = Vec::new();
    for (ordered_index, &command_index) in order.iter().enumerate() {
        let mut barrier_resources = Vec::new();
        for current in &accesses[command_index.index()] {
            let previous = order[..ordered_index]
                .iter()
                .rev()
                .flat_map(|&prior_index| accesses[prior_index.index()].iter().rev())
                .find(|previous| {
                    previous.resource == current.resource
                        && previous.scope.conflicts(&current.scope)
                });
            if let Some(previous) = previous {
                let (old_layout, new_layout) = layouts(previous.scope, current.scope);
                barrier_resources.push(BarrierResource {
                    id: current.resource,
                    range: scope_range(current.scope),
                    src_stage: previous.scope.stage,
                    src_access: previous.scope.access,
                    dst_stage: current.scope.stage,
                    dst_access: current.scope.access,
                    old_layout,
                    new_layout,
                });
                continue;
            }

            let desc = resource_descs[&current.resource];
            let initial = desc.initial_state();
            let (old_layout, new_layout) = layouts_from_state(initial, current.scope);
            let layout_transition = old_layout != new_layout && new_layout.is_some();
            let initial_hazard = matches!(
                desc,
                ResourceDesc::ImportedBuffer { .. } | ResourceDesc::ImportedTexture { .. }
            ) && !initial.access.is_empty()
                && (access_flag_is_write(initial.access)
                    || access_flag_is_write(current.scope.access));
            if layout_transition || initial_hazard {
                barrier_resources.push(BarrierResource {
                    id: current.resource,
                    range: scope_range(current.scope),
                    src_stage: initial.stage,
                    src_access: initial.access,
                    dst_stage: current.scope.stage,
                    dst_access: current.scope.access,
                    old_layout,
                    new_layout,
                });
            }
        }
        if !barrier_resources.is_empty() {
            result.push(PipelineBarrier {
                before_command: ordered_index.into(),
                resources: barrier_resources,
            });
        }
    }
    result
}

fn layouts(
    previous: SubresourceScope,
    current: SubresourceScope,
) -> (Option<vk::ImageLayout>, Option<vk::ImageLayout>) {
    let layout = |scope: SubresourceScope| match scope.subresource {
        SubresourceScopeInner::Texture { layout, .. } => Some(layout),
        SubresourceScopeInner::Buffer(_) => None,
    };
    (layout(previous), layout(current))
}

fn layouts_from_state(
    initial: crate::resources::ResourceState,
    current: SubresourceScope,
) -> (Option<vk::ImageLayout>, Option<vk::ImageLayout>) {
    let current_layout = match current.subresource {
        SubresourceScopeInner::Texture { layout, .. } => Some(layout),
        SubresourceScopeInner::Buffer(_) => None,
    };
    (initial.layout, current_layout)
}

fn bind_resources(
    resources: &BTreeMap<ResourceId, ResourceDesc>,
    accesses: &[Vec<crate::recorder::ResourceAccess>],
) -> BTreeMap<ResourceId, ResourceLifetime> {
    let mut lifetimes = BTreeMap::new();
    for (command, command_accesses) in accesses.iter().enumerate() {
        for access in command_accesses {
            lifetimes
                .entry(access.resource)
                .and_modify(|lifetime: &mut ResourceLifetime| {
                    lifetime.last_command = command.into()
                })
                .or_insert(ResourceLifetime {
                    first_command: command.into(),
                    last_command: command.into(),
                    memory_offset: 0,
                });
        }
    }

    // Skeleton aliasing policy: resources with non-overlapping lifetimes may
    // share a heap offset. Actual size/alignment-aware Vulkan allocation is TODO.
    let mut slots: Vec<(u64, CommandIndex, u64)> = Vec::new(); // (offset, last use, capacity)
    let mut next_offset = 0;
    let mut ids: Vec<_> = lifetimes.keys().copied().collect();
    ids.sort_by_key(|id| lifetimes[id].first_command);
    for id in ids {
        let lifetime = lifetimes.get_mut(&id).expect("lifetime key exists");
        let size = match resources[&id] {
            ResourceDesc::CreatedBuffer(desc) => desc.size,
            ResourceDesc::CreatedTexture(desc) => u64::from(desc.extent.width)
                .saturating_mul(u64::from(desc.extent.height))
                .saturating_mul(u64::from(desc.extent.depth))
                .saturating_mul(4),
            ResourceDesc::ImportedBuffer { .. } | ResourceDesc::ImportedTexture { .. } => 0,
        };
        if resources[&id].is_created()
            && let Some((offset, last_use, capacity)) = slots
                .iter_mut()
                .find(|(_, last, cap)| *last < lifetime.first_command && *cap >= size)
        {
            lifetime.memory_offset = *offset;
            *last_use = lifetime.last_command;
            let _ = capacity;
        } else {
            lifetime.memory_offset = next_offset;
            if resources[&id].is_created() {
                slots.push((next_offset, lifetime.last_command, size));
                next_offset = next_offset.saturating_add(size);
            }
        }
    }
    lifetimes
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{commands::CopyBuffer, resources::BufferRange};

    #[test]
    fn copy_buffer_compiles_with_dependency_and_barrier() {
        let mut graph = RenderGraph::new();
        let src = graph.create_buffer(crate::resources::BufferDesc {
            size: 64,
            usage: vk::BufferUsageFlags::TRANSFER_SRC,
        });
        let dst = graph.create_buffer(crate::resources::BufferDesc {
            size: 64,
            usage: vk::BufferUsageFlags::TRANSFER_DST,
        });
        graph.push(CopyBuffer {
            src,
            src_range: BufferRange {
                offset: 0,
                size: 64,
            },
            dst,
            dst_range: BufferRange {
                offset: 0,
                size: 64,
            },
        });

        let compiled = graph.compile().expect("graph compiles");
        assert_eq!(compiled.command_order(), &[CommandIndex(0)]);
        assert!(compiled.barriers().is_empty());
        assert_eq!(compiled.resource_lifetimes().len(), 2);
    }

    #[test]
    fn write_then_read_creates_order_and_barrier() {
        let mut graph = RenderGraph::new();
        let a = graph.create_buffer(crate::resources::BufferDesc {
            size: 32,
            usage: vk::BufferUsageFlags::TRANSFER_SRC | vk::BufferUsageFlags::TRANSFER_DST,
        });
        let b = graph.create_buffer(crate::resources::BufferDesc {
            size: 32,
            usage: vk::BufferUsageFlags::TRANSFER_DST,
        });
        let c = graph.create_buffer(crate::resources::BufferDesc {
            size: 32,
            usage: vk::BufferUsageFlags::TRANSFER_DST,
        });
        graph.push(CopyBuffer {
            src: a,
            src_range: BufferRange {
                offset: 0,
                size: 32,
            },
            dst: b,
            dst_range: BufferRange {
                offset: 0,
                size: 32,
            },
        });
        graph.push(CopyBuffer {
            src: b,
            src_range: BufferRange {
                offset: 0,
                size: 32,
            },
            dst: c,
            dst_range: BufferRange {
                offset: 0,
                size: 32,
            },
        });

        let compiled = graph.compile().expect("graph compiles");
        assert_eq!(
            compiled.command_order(),
            &[CommandIndex(0), CommandIndex(1)]
        );
        assert_eq!(compiled.barriers().len(), 1);
        assert_eq!(compiled.barriers()[0].before_command, CommandIndex(1));
        assert_eq!(compiled.barriers()[0].resources[0].id, b.id());
    }
}