monolith

use crate::image::{ImageInner, color_range};

use ash::vk;

use std::sync::Arc;

/// Layout changes recorded into a command buffer that has not been submitted.

/// Images keep their old layout until `commit`, so a failed submit leaves the

/// tracking true to what the GPU has done.

#[derive(Default)]

pub(crate) struct LayoutLog(Vec<(Arc<ImageInner>, vk::ImageLayout)>);

impl LayoutLog {

    fn current(&self, image: &Arc<ImageInner>) -> vk::ImageLayout {

        self.0

            .iter()

            .rev()

            .find(|(i, _)| Arc::ptr_eq(i, image))

            .map_or_else(|| image.layout(), |(_, layout)| *layout)

    }

    pub fn commit(self) {

        for (image, layout) in self.0 {

            image.set_layout(layout);

        }

    }

}

struct Scope {

    stage: vk::PipelineStageFlags2,

    access: vk::AccessFlags2,

}

/// What an image may still be doing when it leaves `layout`. Layouts the

/// renderer does not use itself, the general layout of dmabufs among them,

/// wait for everything.

fn leaving(layout: vk::ImageLayout) -> Scope {

    let (stage, access) = match layout {

        vk::ImageLayout::UNDEFINED => (vk::PipelineStageFlags2::NONE, vk::AccessFlags2::NONE),

        vk::ImageLayout::TRANSFER_DST_OPTIMAL => (

            vk::PipelineStageFlags2::TRANSFER,

            vk::AccessFlags2::TRANSFER_WRITE,

        ),

        vk::ImageLayout::TRANSFER_SRC_OPTIMAL => {

            (vk::PipelineStageFlags2::TRANSFER, vk::AccessFlags2::NONE)

        }

        vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL => (

            vk::PipelineStageFlags2::FRAGMENT_SHADER,

            vk::AccessFlags2::NONE,

        ),

        vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL => (

            vk::PipelineStageFlags2::COLOR_ATTACHMENT_OUTPUT,

            vk::AccessFlags2::COLOR_ATTACHMENT_WRITE,

        ),

        _ => (

            vk::PipelineStageFlags2::ALL_COMMANDS,

            vk::AccessFlags2::MEMORY_WRITE,

        ),

    };

    Scope { stage, access }

}

/// What an image does once it is in `layout`.

fn entering(layout: vk::ImageLayout) -> Scope {

    let (stage, access) = match layout {

        vk::ImageLayout::TRANSFER_DST_OPTIMAL => (

            vk::PipelineStageFlags2::TRANSFER,

            vk::AccessFlags2::TRANSFER_WRITE,

        ),

        vk::ImageLayout::TRANSFER_SRC_OPTIMAL => (

            vk::PipelineStageFlags2::TRANSFER,

            vk::AccessFlags2::TRANSFER_READ,

        ),

        vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL => (

            vk::PipelineStageFlags2::FRAGMENT_SHADER,

            vk::AccessFlags2::SHADER_SAMPLED_READ,

        ),

        vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL => (

            vk::PipelineStageFlags2::COLOR_ATTACHMENT_OUTPUT,

            vk::AccessFlags2::COLOR_ATTACHMENT_READ | vk::AccessFlags2::COLOR_ATTACHMENT_WRITE,

        ),

        _ => (

            vk::PipelineStageFlags2::ALL_COMMANDS,

            vk::AccessFlags2::MEMORY_READ | vk::AccessFlags2::MEMORY_WRITE,

        ),

    };

    Scope { stage, access }

}

/// The image barriers of one phase of a command buffer, recorded together.

#[derive(Default)]

pub(crate) struct Barriers(Vec<vk::ImageMemoryBarrier2<'static>>);

impl Barriers {

    fn push(

        &mut self,

        image: vk::Image,

        (old, src_family): (vk::ImageLayout, u32),

        (new, dst_family): (vk::ImageLayout, u32),

    ) {

        let (from, to) = (leaving(old), entering(new));

        self.0.push(

            vk::ImageMemoryBarrier2::default()

                .src_stage_mask(from.stage)

                .src_access_mask(from.access)

                .dst_stage_mask(to.stage)

                .dst_access_mask(to.access)

                .old_layout(old)

                .new_layout(new)

                .src_queue_family_index(src_family)

                .dst_queue_family_index(dst_family)

                .image(image)

                .subresource_range(color_range()),

        );

    }

    /// Moves an image into `layout` for use on this queue. Dmabuf images come

    /// from the foreign queue family in the general layout, others continue

    /// from their tracked layout.

    pub fn acquire(

        &mut self,

        family: u32,

        image: &Arc<ImageInner>,

        layout: vk::ImageLayout,

        log: &mut LayoutLog,

    ) {

        if image.external {

            self.push(

                image.image,

                (vk::ImageLayout::GENERAL, vk::QUEUE_FAMILY_FOREIGN_EXT),

                (layout, family),

            );

        } else {

            let old = log.current(image);

            if old != layout {

                self.push(

                    image.image,

                    (old, vk::QUEUE_FAMILY_IGNORED),

                    (layout, vk::QUEUE_FAMILY_IGNORED),

                );

                log.0.push((image.clone(), layout));

            }

        }

    }

    /// Hands a dmabuf image back to the foreign queue family after use.

    pub fn release(&mut self, family: u32, image: &ImageInner, layout: vk::ImageLayout) {

        if image.external {

            self.push(

                image.image,

                (layout, family),

                (vk::ImageLayout::GENERAL, vk::QUEUE_FAMILY_FOREIGN_EXT),

            );

        }

    }

    pub unsafe fn record(self, d: &ash::Device, cmd: vk::CommandBuffer) {

        if self.0.is_empty() {

            return;

        }

        unsafe {

            d.cmd_pipeline_barrier2(

                cmd,

                &vk::DependencyInfo::default().image_memory_barriers(&self.0),

            )

        };

    }

}