monolith

use crate::Error;

use crate::device::Shared;

use crate::dmabuf::{self, DmabufCaps, RENDER_USAGE, SAMPLE_USAGE, VulkanAllocator};

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::allocator::dmabuf::{Dmabuf, WeakDmabuf};

use smithay::backend::allocator::format::FormatSet;

use smithay::backend::renderer::ImportMemWl;

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

use smithay::backend::renderer::{

    Bind, ContextId, DebugFlags, ExportMem, ImportDma, ImportDmaWl, ImportMem, Offscreen, Renderer,

    RendererSuper, Texture, TextureFilter,

};

use smithay::reexports::wayland_server::Resource;

use smithay::reexports::wayland_server::backend::ObjectId;

use smithay::reexports::wayland_server::protocol::wl_buffer::WlBuffer;

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

use smithay::wayland::compositor::SurfaceData;

use smithay::wayland::shm::{shm_format_to_fourcc, with_buffer_contents};

use std::cell::Cell;

use std::collections::HashMap;

use std::fmt;

use std::marker::PhantomData;

use std::sync::{Arc, Mutex};

use std::time::{Duration, Instant};

struct Pending {

    fence: Arc<FenceInner>,

    cmd: vk::CommandBuffer,

    _keep: Vec<Arc<ImageInner>>,

    _staging: Option<Staging>,

}

pub struct VulkanRenderer {

    pending: Vec<Pending>,

    dmabuf_textures: HashMap<WeakDmabuf, VulkanTexture>,

    dmabuf_targets: HashMap<WeakDmabuf, Arc<ImageInner>>,

    pipelines: Pipelines,

    pool: vk::CommandPool,

    /// Command buffers whose work finished, reset and reused by `begin`.

    spare_cmds: Vec<vk::CommandBuffer>,

    samplers: [vk::Sampler; 4],

    min_filter: TextureFilter,

    mag_filter: TextureFilter,

    debug: DebugFlags,

    context: ContextId<VulkanTexture>,

    caps: Arc<DmabufCaps>,

    shared: Arc<Shared>,

    submissions: Cell<u64>,

    /// Shm uploads staged off the main thread, by buffer.

    staged: HashMap<ObjectId, (WlBuffer, StagedShm)>,

    stall: Cell<StallPhase>,

}

/// Where in every 1024 submissions the driver stalls. A phase is trusted

/// once two stalls land on it.

#[derive(Clone, Copy, Default)]

struct StallPhase {

    phase: Option<u64>,

    seen: Option<u64>,

}

impl StallPhase {

    fn saw(self, at: u64) -> Self {

        if self.seen == Some(at) {

            if self.phase != Some(at) {

                tracing::debug!(phase = at, "queue submissions stall periodically");

            }

            Self {

                phase: Some(at),

                seen: Some(at),

            }

        } else {

            Self {

                seen: Some(at),

                ..self

            }

        }

    }

}

/// A submission this slow is a driver stall rather than work.

const STALL: Duration = Duration::from_micros(1500);

/// How many submissions early a stall is absorbed. Each empty one costs

/// tens of microseconds.

const ABSORB_AHEAD: u64 = 8;

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> {

        Self::with_node(None)

    }

    /// A renderer on the GPU behind DRM device number `node`, so it can

    /// render into buffers that GPU scans out.

    pub fn for_node(node: u64) -> Result<Self, Error> {

        Self::with_node(Some(node))

    }

    fn with_node(node: Option<u64>) -> Result<Self, Error> {

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

        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

                            | vk::CommandPoolCreateFlags::RESET_COMMAND_BUFFER,

                    ),

                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(),

            dmabuf_textures: HashMap::new(),

            dmabuf_targets: HashMap::new(),

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

            pool,

            spare_cmds: Vec::new(),

            submissions: Cell::new(0),

            staged: HashMap::new(),

            stall: Cell::new(StallPhase {

                // Measured on an RTX 3060, see research/nvidia.md.

                phase: shared.device.name.contains("NVIDIA").then_some(1015),

                seen: None,

            }),

            samplers,

            min_filter: TextureFilter::Linear,

            mag_filter: TextureFilter::Linear,

            debug: DebugFlags::empty(),

            context: ContextId::new(),

            caps: Arc::new(DmabufCaps::query(&shared)),

            shared,

        })

    }

    /// Name of the physical device in use.

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

        &self.shared.device.name

    }

    /// Whether dmabufs can be imported, rendered into and allocated.

    pub fn supports_dmabuf(&self) -> bool {

        self.caps.render.iter().next().is_some()

    }

    /// Formats and modifiers this renderer can render into as dmabufs.

    pub fn render_formats(&self) -> FormatSet {

        self.caps.render.clone()

    }

    /// An allocator for render-capable dmabufs on this device.

    pub fn allocator(&self) -> VulkanAllocator {

        VulkanAllocator {

            shared: self.shared.clone(),

            caps: self.caps.clone(),

        }

    }

    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 spare = &mut self.spare_cmds;

        self.pending.retain(|p| {

            if !p.fence.signaled() {

                return true;

            }

            spare.push(p.cmd);

            false

        });

        self.dmabuf_textures.retain(|k, _| !k.is_gone());

        self.dmabuf_targets.retain(|k, _| !k.is_gone());

    }

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

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

        let cmd = match self.spare_cmds.pop() {

            Some(cmd) => {

                unsafe { d.reset_command_buffer(cmd, vk::CommandBufferResetFlags::empty())? };

                cmd

            }

            None => 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)

    }

    #[profiling::function]

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

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

        {

            profiling::scope!("end_command_buffer");

            unsafe { d.end_command_buffer(cmd)? };

        }

        let fence = {

            profiling::scope!("create_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());

        let began = Instant::now();

        {

            profiling::scope!("queue_submit");

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

        }

        let n = self.submissions.get();

        self.submissions.set(n + 1);

        if began.elapsed() > STALL {

            self.stall.set(self.stall.get().saw(n % 1024));

        }

        Ok(fence)

    }

    /// Spends an upcoming stalling submission on empty command buffers, so

    /// the stall lands now, in idle time, instead of in a frame. Call it

    /// right after a flip. NVIDIA stalls for milliseconds on every 1024th

    /// submission, see research/nvidia.md.

    pub fn absorb_submit_stall(&mut self) -> Result<(), Error> {

        self.reap();

        let Some(phase) = self.stall.get().phase else {

            return Ok(());

        };

        let ahead = (phase + 1024 - self.submissions.get() % 1024) % 1024;

        if ahead >= ABSORB_AHEAD {

            return Ok(());

        }

        for _ in 0..=ahead {

            let cmd = self.begin()?;

            let fence = self.end_and_submit(cmd)?;

            self.pending.push(Pending {

                fence,

                cmd,

                _keep: Vec::new(),

                _staging: None,

            });

        }

        Ok(())

    }

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

    fn one_shot(

        &mut self,

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

    ) -> Result<(), Error> {

        let cmd = self.begin()?;

        record(

            &self.shared.device.device,

            cmd,

            self.shared.device.queue_family,

        );

        let fence = self.end_and_submit(cmd)?;

        let res = fence.wait();

        self.spare_cmds.push(cmd);

        Ok(res?)

    }

    #[profiling::function]

    pub(crate) fn submit(

        &mut self,

        target: &mut VulkanTarget<'_>,

        ops: Vec<Op>,

    ) -> Result<SyncPoint, Error> {

        self.reap();

        let image = target.0.clone();

        let mut sampled: Vec<Arc<ImageInner>> = Vec::new();

        for op in &ops {

            if let Op::Draw {

                texture: Some(t), ..

            } = op

                && !sampled.iter().any(|s| Arc::ptr_eq(s, &t.0))

            {

                sampled.push(t.0.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 family = self.shared.device.queue_family;

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

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

        let sampler = self.sampler();

        unsafe {

            for t in &sampled {

                acquire(d, cmd, family, t, vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL);

            }

            acquire(

                d,

                cmd,

                family,

                &image,

                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);

            release(

                d,

                cmd,

                family,

                &image,

                vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL,

            );

            for t in &sampled {

                release(d, cmd, family, t, vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL);

            }

        }

        let fence = self.end_and_submit(cmd)?;

        sampled.push(image);

        self.pending.push(Pending {

            fence: fence.clone(),

            cmd,

            _keep: sampled,

            _staging: None,

        });

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