monolith

use crate::Error;

use crate::device::Shared;

use crate::format::{BYTES_PER_PIXEL, SUPPORTED};

use crate::frame::{Op, VulkanFrame};

use crate::image::{

    ImageInner, Staging, VulkanImage, VulkanMapping, VulkanTarget, VulkanTexture, color_range,

};

use crate::pipeline::Pipelines;

use crate::sync::{FenceInner, VkFence};

use ash::vk;

use gpu_allocator::MemoryLocation;

use smithay::backend::allocator::Fourcc;

use smithay::backend::renderer::sync::SyncPoint;

use smithay::backend::renderer::{

    Bind, ContextId, DebugFlags, ExportMem, ImportMem, Offscreen, Renderer, RendererSuper, Texture,

    TextureFilter,

};

use smithay::utils::{Buffer, Physical, Rectangle, Size, Transform};

use std::fmt;

use std::sync::Arc;

struct Pending {

    fence: Arc<FenceInner>,

    cmd: vk::CommandBuffer,

    _keep: Vec<VulkanTexture>,

}

pub struct VulkanRenderer {

    pending: Vec<Pending>,

    pipelines: Pipelines,

    pool: vk::CommandPool,

    samplers: [vk::Sampler; 4],

    min_filter: TextureFilter,

    mag_filter: TextureFilter,

    debug: DebugFlags,

    context: ContextId<VulkanTexture>,

    shared: Arc<Shared>,

}

impl fmt::Debug for VulkanRenderer {

    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {

        f.debug_struct("VulkanRenderer")

            .field("device", &self.shared.device.name)

            .finish()

    }

}

impl VulkanRenderer {

    pub fn new() -> Result<Self, Error> {

        let shared = Arc::new(Shared::new()?);

        let d = &shared.device.device;

        let pool = unsafe {

            d.create_command_pool(

                &vk::CommandPoolCreateInfo::default()

                    .queue_family_index(shared.device.queue_family)

                    .flags(vk::CommandPoolCreateFlags::TRANSIENT),

                None,

            )?

        };

        let filter = |f: TextureFilter| match f {

            TextureFilter::Linear => vk::Filter::LINEAR,

            TextureFilter::Nearest => vk::Filter::NEAREST,

        };

        let mut samplers = [vk::Sampler::null(); 4];

        for (i, s) in samplers.iter_mut().enumerate() {

            let (min, mag) = (sampler_filter(i >> 1), sampler_filter(i & 1));

            *s = unsafe {

                d.create_sampler(

                    &vk::SamplerCreateInfo::default()

                        .min_filter(filter(min))

                        .mag_filter(filter(mag))

                        .address_mode_u(vk::SamplerAddressMode::CLAMP_TO_EDGE)

                        .address_mode_v(vk::SamplerAddressMode::CLAMP_TO_EDGE)

                        .address_mode_w(vk::SamplerAddressMode::CLAMP_TO_EDGE),

                    None,

                )?

            };

        }

        Ok(Self {

            pending: Vec::new(),

            pipelines: Pipelines::new(&shared)?,

            pool,

            samplers,

            min_filter: TextureFilter::Linear,

            mag_filter: TextureFilter::Linear,

            debug: DebugFlags::empty(),

            context: ContextId::new(),

            shared,

        })

    }

    /// Name of the physical device in use.

    pub fn device_name(&self) -> &str {

        &self.shared.device.name

    }

    pub(crate) fn shared(&self) -> &Arc<Shared> {

        &self.shared

    }

    pub(crate) fn context_id_ref(&self) -> &ContextId<VulkanTexture> {

        &self.context

    }

    fn sampler(&self) -> vk::Sampler {

        let idx = |f: TextureFilter| usize::from(f == TextureFilter::Nearest);

        self.samplers[(idx(self.min_filter) << 1) | idx(self.mag_filter)]

    }

    fn reap(&mut self) {

        let d = &self.shared.device.device;

        let pool = self.pool;

        self.pending.retain(|p| {

            if !p.fence.signaled() {

                return true;

            }

            unsafe { d.free_command_buffers(pool, &[p.cmd]) };

            false

        });

    }

    fn begin(&self) -> Result<vk::CommandBuffer, Error> {

        let d = &self.shared.device.device;

        let cmd = unsafe {

            d.allocate_command_buffers(

                &vk::CommandBufferAllocateInfo::default()

                    .command_pool(self.pool)

                    .level(vk::CommandBufferLevel::PRIMARY)

                    .command_buffer_count(1),

            )?[0]

        };

        unsafe {

            d.begin_command_buffer(

                cmd,

                &vk::CommandBufferBeginInfo::default()

                    .flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT),

            )?

        };

        Ok(cmd)

    }

    fn end_and_submit(&self, cmd: vk::CommandBuffer) -> Result<Arc<FenceInner>, Error> {

        let d = &self.shared.device.device;

        unsafe { d.end_command_buffer(cmd)? };

        let fence = FenceInner::new(&self.shared)?;

        let cmds = [vk::CommandBufferSubmitInfo::default().command_buffer(cmd)];

        let submit = vk::SubmitInfo2::default().command_buffer_infos(&cmds);

        let queue = self

            .shared

            .device

            .queue

            .lock()

            .unwrap_or_else(|e| e.into_inner());

        unsafe { d.queue_submit2(*queue, &[submit], fence.fence)? };

        Ok(fence)

    }

    /// Records, submits and waits for a short command buffer.

    fn one_shot(

        &mut self,

        record: impl FnOnce(&ash::Device, vk::CommandBuffer),

    ) -> Result<(), Error> {

        let cmd = self.begin()?;

        record(&self.shared.device.device, cmd);

        let fence = self.end_and_submit(cmd)?;

        let res = fence.wait();

        unsafe {

            self.shared

                .device

                .device

                .free_command_buffers(self.pool, &[cmd])

        };

        Ok(res?)

    }

    pub(crate) fn submit(

        &mut self,

        target: &mut VulkanTarget<'_>,

        ops: Vec<Op>,

    ) -> Result<SyncPoint, Error> {

        self.reap();

        let image = target.0.0.clone();

        let mut keep = vec![VulkanTexture(image.clone())];

        let mut transitions = Vec::new();

        for op in &ops {

            if let Op::Draw {

                texture: Some(t), ..

            } = op

            {

                if t.0.layout() != vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL {

                    transitions.push(t.clone());

                }

                keep.push(t.clone());

            }

        }

        let mut pipelines = Vec::with_capacity(ops.len());

        for op in &ops {

            pipelines.push(match op {

                Op::Draw { kind, .. } => Some(self.pipelines.get(image.format, *kind)?),

                Op::Clear { .. } => None,

            });

        }

        let cmd = self.begin()?;

        let d = &self.shared.device.device;

        let push = &self.shared.device.push_descriptor;

        let (w, h) = (image.width, image.height);

        let sampler = self.sampler();

        unsafe {

            for t in &transitions {

                barrier(

                    d,

                    cmd,

                    t.0.image,

                    t.0.layout(),

                    vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL,

                );

                t.0.set_layout(vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL);

            }

            barrier(

                d,

                cmd,

                image.image,

                image.layout(),

                vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL,

            );

            let attachments = [vk::RenderingAttachmentInfo::default()

                .image_view(image.view)

                .image_layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL)

                .load_op(vk::AttachmentLoadOp::LOAD)

                .store_op(vk::AttachmentStoreOp::STORE)];

            let area = vk::Rect2D {

                offset: vk::Offset2D::default(),

                extent: vk::Extent2D {

                    width: w,

                    height: h,

                },

            };

            d.cmd_begin_rendering(

                cmd,

                &vk::RenderingInfo::default()

                    .render_area(area)

                    .layer_count(1)

                    .color_attachments(&attachments),

            );

            d.cmd_set_viewport(

                cmd,

                0,

                &[vk::Viewport {

                    x: 0.0,

                    y: 0.0,

                    width: w as f32,

                    height: h as f32,

                    min_depth: 0.0,

                    max_depth: 1.0,

                }],

            );

            let mut bound = vk::Pipeline::null();

            for (op, pipeline) in ops.iter().zip(pipelines) {

                match op {

                    Op::Clear { color, rects } => {

                        let clear = [vk::ClearAttachment {

                            aspect_mask: vk::ImageAspectFlags::COLOR,

                            color_attachment: 0,

                            clear_value: vk::ClearValue {

                                color: vk::ClearColorValue { float32: *color },

                            },

                        }];

                        let rects: Vec<_> = rects

                            .iter()

                            .map(|&rect| vk::ClearRect {

                                rect,

                                base_array_layer: 0,

                                layer_count: 1,

                            })

                            .collect();

                        d.cmd_clear_attachments(cmd, &clear, &rects);

                    }

                    Op::Draw {

                        texture,

                        push: constants,

                        scissors,

                        ..

                    } => {

                        let pipeline = pipeline.unwrap_or_default();

                        if pipeline != bound {

                            d.cmd_bind_pipeline(cmd, vk::PipelineBindPoint::GRAPHICS, pipeline);

                            bound = pipeline;

                        }

                        if let Some(t) = texture {

                            let info = [vk::DescriptorImageInfo::default()

                                .image_view(t.0.view)

                                .image_layout(vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL)

                                .sampler(sampler)];

                            let write = vk::WriteDescriptorSet::default()

                                .dst_binding(0)

                                .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)

                                .image_info(&info);

                            push.cmd_push_descriptor_set(

                                cmd,

                                vk::PipelineBindPoint::GRAPHICS,

                                self.pipelines.layout,

                                0,

                                &[write],

                            );

                        }

                        d.cmd_push_constants(

                            cmd,

                            self.pipelines.layout,

                            vk::ShaderStageFlags::VERTEX | vk::ShaderStageFlags::FRAGMENT,

                            0,

                            constants.bytes(),

                        );

                        for s in scissors {

                            d.cmd_set_scissor(cmd, 0, &[*s]);

                            d.cmd_draw(cmd, 4, 1, 0, 0);

                        }

                    }

                }

            }

            d.cmd_end_rendering(cmd);

        }

        image.set_layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL);

        let fence = self.end_and_submit(cmd)?;

        self.pending.push(Pending {

            fence: fence.clone(),

            cmd,

            _keep: keep,

        });

        Ok(SyncPoint::from(VkFence(fence)))

    }

    fn upload(

        &mut self,

        image: &Arc<ImageInner>,

        data: &[u8],

        region: Rectangle<i32, Buffer>,

        flipped: bool,

    ) -> Result<(), Error> {

        let width = image.width as usize;

        let needed = width * image.height as usize * BYTES_PER_PIXEL;

        if data.len() < needed {

            return Err(Error::BufferTooSmall {

                needed,

                got: data.len(),

            });

        }

        let mut staging = Staging::new(

            &self.shared,

            needed,

            vk::BufferUsageFlags::TRANSFER_SRC,

            MemoryLocation::CpuToGpu,

        )?;

        let row = width * BYTES_PER_PIXEL;

        let rows = image.height as usize;

        let dst = staging.bytes_mut();

        for y in 0..rows {

            let from = if flipped { rows - 1 - y } else { y };

            dst[y * row..(y + 1) * row].copy_from_slice(&data[from * row..(from + 1) * row]);

        }

        let old = image.layout();

        let copy = vk::BufferImageCopy::default()

            .buffer_offset(

                ((region.loc.y as usize * width + region.loc.x as usize) * BYTES_PER_PIXEL) as u64,

            )

            .buffer_row_length(image.width)

            .image_subresource(

                vk::ImageSubresourceLayers::default()

                    .aspect_mask(vk::ImageAspectFlags::COLOR)

                    .layer_count(1),

            )

            .image_offset(vk::Offset3D {

                x: region.loc.x,

                y: region.loc.y,

                z: 0,

            })

            .image_extent(vk::Extent3D {

                width: region.size.w as u32,

                height: region.size.h as u32,

                depth: 1,

            });

        let (img, buffer) = (image.image, staging.buffer);

        self.one_shot(|d, cmd| unsafe {

            barrier(d, cmd, img, old, vk::ImageLayout::TRANSFER_DST_OPTIMAL);

            d.cmd_copy_buffer_to_image(

                cmd,

                buffer,

                img,

                vk::ImageLayout::TRANSFER_DST_OPTIMAL,

                &[copy],

            );

            barrier(

                d,

                cmd,

                img,

                vk::ImageLayout::TRANSFER_DST_OPTIMAL,

                vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL,

            );

        })?;

        image.set_layout(vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL);

        Ok(())

    }

    fn read_back(

        &mut self,

        image: &Arc<ImageInner>,

        region: Rectangle<i32, Buffer>,

        fourcc: Fourcc,

    ) -> Result<VulkanMapping, Error> {

        let region = region

            .intersection(Rectangle::from_size(

                (image.width as i32, image.height as i32).into(),

            ))

            .unwrap_or_default();

        let (w, h) = (region.size.w.max(0) as u32, region.size.h.max(0) as u32);

        let swap = needs_swizzle(image.fourcc, fourcc)?;

        let len = (w as usize * h as usize * BYTES_PER_PIXEL).max(BYTES_PER_PIXEL);

        let staging = Staging::new(

            &self.shared,

            len,

            vk::BufferUsageFlags::TRANSFER_DST,

            MemoryLocation::GpuToCpu,

        )?;

        if w > 0 && h > 0 {

            let old = image.layout();

            let copy = vk::BufferImageCopy::default()

                .image_subresource(

                    vk::ImageSubresourceLayers::default()

                        .aspect_mask(vk::ImageAspectFlags::COLOR)

                        .layer_count(1),

                )

                .image_offset(vk::Offset3D {

                    x: region.loc.x,

                    y: region.loc.y,

                    z: 0,

                })

                .image_extent(vk::Extent3D {

                    width: w,

                    height: h,

                    depth: 1,

                });

            let (img, buffer) = (image.image, staging.buffer);

            self.one_shot(|d, cmd| unsafe {

                barrier(d, cmd, img, old, vk::ImageLayout::TRANSFER_SRC_OPTIMAL);

                d.cmd_copy_image_to_buffer(

                    cmd,

                    img,

                    vk::ImageLayout::TRANSFER_SRC_OPTIMAL,

                    buffer,

                    &[copy],

                );

            })?;

            image.set_layout(vk::ImageLayout::TRANSFER_SRC_OPTIMAL);

        }

        let mut data = staging.bytes()[..w as usize * h as usize * BYTES_PER_PIXEL].to_vec();

        if swap {

            for px in data.as_chunks_mut::<BYTES_PER_PIXEL>().0 {

                px.swap(0, 2);

            }

        }

        Ok(VulkanMapping {

            data,

            width: w,

            height: h,

            fourcc,

        })

    }

    fn own(&self, image: &Arc<ImageInner>) -> Result<(), Error> {

        if Arc::ptr_eq(&image.shared, &self.shared) {

            Ok(())

        } else {

            Err(Error::ForeignTexture)

        }

    }

}

impl Drop for VulkanRenderer {