monolith

use crate::Error;

use ash::vk;

use gpu_allocator::{AllocationSizes, AllocatorDebugSettings};

use gpu_allocator::vulkan::{Allocator, AllocatorCreateDesc};

use std::mem::ManuallyDrop;

use std::sync::Mutex;

/// Everything that must outlive every Vulkan object the renderer hands out.

/// Fields drop in declaration order, so the allocator frees its memory before

/// the device is destroyed.

pub(crate) struct Shared {

    pub allocator: ManuallyDrop<Mutex<Allocator>>,

    pub device: Device,

}

impl Drop for Shared {

    fn drop(&mut self) {

        if !crate::error::device_lost() {

            unsafe { ManuallyDrop::drop(&mut self.allocator) };

        }

    }

}

/// Loaders for the extensions that dmabuf import, export and fence export

/// need. Absent when the device lacks any of them.

pub(crate) struct DmabufExt {

    pub memory_fd: ash::khr::external_memory_fd::Device,

    pub fence_fd: ash::khr::external_fence_fd::Device,

    pub modifier: ash::ext::image_drm_format_modifier::Device,

}

const DMABUF_EXTENSIONS: [&std::ffi::CStr; 5] = [

    ash::khr::external_memory_fd::NAME,

    ash::ext::external_memory_dma_buf::NAME,

    ash::ext::image_drm_format_modifier::NAME,

    ash::ext::queue_family_foreign::NAME,

    ash::khr::external_fence_fd::NAME,

];

pub(crate) struct Device {

    pub _entry: ash::Entry,

    pub instance: ash::Instance,

    pub physical: vk::PhysicalDevice,

    pub device: ash::Device,

    pub queue: Mutex<vk::Queue>,

    pub queue_family: u32,

    pub push_descriptor: ash::khr::push_descriptor::Device,

    pub dmabuf: Option<DmabufExt>,

    pub name: String,

    pub drm_node: Option<u64>,

    /// Signalled fences kept for reuse, see `FenceInner`.

    pub spare_fences: Mutex<Vec<vk::Fence>>,

}

/// A lost or hung device is leaked. Waiting on it or destroying it can block

/// forever in the driver, and the process is about to exit anyway.

impl Drop for Device {

    fn drop(&mut self) {

        if crate::error::device_lost() {

            return;

        }

        unsafe {

            let _ = self.device.device_wait_idle();

            for f in self

                .spare_fences

                .get_mut()

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

                .drain(..)

            {

                self.device.destroy_fence(f, None);

            }

            self.device.destroy_device(None);

            self.instance.destroy_instance(None);

        }

    }

}

impl Shared {

    /// Opens the first suitable device, or the one whose name contains

    /// `MONOLITH_VK_DEVICE` when that is set. With `node`, only the device

    /// behind that DRM device number qualifies.

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

        let entry = unsafe { ash::Entry::load()? };

        let app = vk::ApplicationInfo::default()

            .application_name(c"monolith")

            .engine_name(c"monolith")

            .api_version(vk::API_VERSION_1_3);

        let wanted = std::env::var("MONOLITH_VK_DEVICE").ok();

        let mut env = node.map(driver_env).unwrap_or_default();

        let (instance, physical, queue_family, name) = loop {

            for (var, value) in &env {

                // SAFETY: the compositor creates its renderer before starting any

                // other thread, so nothing reads the environment concurrently.

                unsafe { std::env::set_var(var, value) };

            }

            let instance = unsafe {

                entry.create_instance(

                    &vk::InstanceCreateInfo::default().application_info(&app),

                    None,

                )

            };

            for (var, _) in &env {

                // SAFETY: as above. Removing them keeps clients from inheriting them.

                unsafe { std::env::remove_var(var) };

            }

            let instance = match instance {

                Ok(i) => i,

                Err(_) if !std::mem::take(&mut env).is_empty() => continue,

                Err(e) => return Err(e.into()),

            };

            match unsafe { pick_physical(&instance, wanted.as_deref(), node) } {

                Some((p, f, n)) => break (instance, p, f, n),

                None => {

                    unsafe { instance.destroy_instance(None) };

                    if std::mem::take(&mut env).is_empty() {

                        return Err(Error::NoSuitableDevice);

                    }

                }

            }

        };

        let priorities = [1.0];

        let available: Vec<_> = unsafe { instance.enumerate_device_extension_properties(physical) }

            .unwrap_or_default()

            .iter()

            .filter_map(|e| e.extension_name_as_c_str().ok().map(ToOwned::to_owned))

            .collect();

        let drm_node = available

            .iter()

            .any(|a| a.as_c_str() == ash::ext::physical_device_drm::NAME)

            .then(|| unsafe { drm_node(&instance, physical) })

            .flatten();

        let has_dmabuf = DMABUF_EXTENSIONS

            .iter()

            .all(|n| available.iter().any(|a| a.as_c_str() == *n));

        let mut extensions = vec![ash::khr::push_descriptor::NAME.as_ptr()];

        if has_dmabuf {

            extensions.extend(DMABUF_EXTENSIONS.iter().map(|n| n.as_ptr()));

        }

        let global_priority = [

            ash::khr::global_priority::NAME,

            ash::ext::global_priority::NAME,

        ]

        .into_iter()

        .find(|n| available.iter().any(|a| a.as_c_str() == *n));

        let mut priorities_to_try = vec![None];

        if let Some(name) = global_priority {

            extensions.push(name.as_ptr());

            priorities_to_try = vec![

                Some(vk::QueueGlobalPriorityKHR::REALTIME),

                Some(vk::QueueGlobalPriorityKHR::HIGH),

                None,

            ];

        }

        let mut created = Err(vk::Result::ERROR_INITIALIZATION_FAILED);

        for priority in priorities_to_try {

            let mut features13 = vk::PhysicalDeviceVulkan13Features::default()

                .dynamic_rendering(true)

                .synchronization2(true);

            let mut global = vk::DeviceQueueGlobalPriorityCreateInfoKHR::default()

                .global_priority(priority.unwrap_or(vk::QueueGlobalPriorityKHR::MEDIUM));

            let mut queue = vk::DeviceQueueCreateInfo::default()

                .queue_family_index(queue_family)

                .queue_priorities(&priorities);

            if priority.is_some() {

                queue = queue.push_next(&mut global);

            }

            let queues = [queue];

            let info = vk::DeviceCreateInfo::default()

                .queue_create_infos(&queues)

                .enabled_extension_names(&extensions)

                .push_next(&mut features13);

            created = unsafe { instance.create_device(physical, &info, None) };

            match &created {

                Ok(_) => {

                    tracing::info!(?priority, "queue global priority");

                    break;

                }

                Err(e) => tracing::debug!(?priority, "queue global priority refused: {e}"),

            }

        }

        let device = match created {

            Ok(d) => d,

            Err(e) => {

                unsafe { instance.destroy_instance(None) };

                return Err(e.into());

            }

        };

        let queue = unsafe { device.get_device_queue(queue_family, 0) };

        let push_descriptor = ash::khr::push_descriptor::Device::new(&instance, &device);

        let dmabuf = has_dmabuf.then(|| DmabufExt {

            memory_fd: ash::khr::external_memory_fd::Device::new(&instance, &device),

            fence_fd: ash::khr::external_fence_fd::Device::new(&instance, &device),

            modifier: ash::ext::image_drm_format_modifier::Device::new(&instance, &device),

        });

        let allocator = Allocator::new(&AllocatorCreateDesc {

            instance: instance.clone(),

            device: device.clone(),

            physical_device: physical,

            // The live compositor is non-dumpable, so with

            // RUST_LOG=gpu_allocator=debug its log is the only way to see

            // what fills the memory it maps.

            debug_settings: AllocatorDebugSettings {

                log_memory_information: true,

                log_allocations: true,

                log_frees: true,

                ..Default::default()

            },

            buffer_device_address: false,

            // Host blocks are pinned system RAM that stays resident while any

            // allocation lives in them, and the default 64 MB block is mostly idle.

            allocation_sizes: AllocationSizes::new(64 << 20, 4 << 20),

        });

        let device = Device {

            _entry: entry,

            instance,

            physical,

            device,

            queue: Mutex::new(queue),

            queue_family,

            push_descriptor,

            dmabuf,

            name,

            drm_node,

            spare_fences: Mutex::new(Vec::new()),

        };

        Ok(Self {

            allocator: ManuallyDrop::new(Mutex::new(allocator?)),

            device,

        })

    }

}

/// Whether `physical` is the device behind DRM device number `node`, as its

/// primary or render node.

unsafe fn serves_node(instance: &ash::Instance, physical: vk::PhysicalDevice, node: u64) -> bool {

    let mut drm = vk::PhysicalDeviceDrmPropertiesEXT::default();

    let mut props = vk::PhysicalDeviceProperties2::default().push_next(&mut drm);

    unsafe { instance.get_physical_device_properties2(physical, &mut props) };

    let (major, minor) = (

        rustix::fs::major(node) as i64,

        rustix::fs::minor(node) as i64,

    );

    (drm.has_primary == vk::TRUE && (drm.primary_major, drm.primary_minor) == (major, minor))

        || (drm.has_render == vk::TRUE && (drm.render_major, drm.render_minor) == (major, minor))

}

/// Variables set only while the instance is created, unless the user already

/// picks drivers themselves. The loader keeps every ICD it finds loaded for the

/// instance's lifetime, on NixOS including lavapipe and 26 MB of resident

/// libLLVM. Without an X server NVIDIA's ICD also initialises through glvnd's

/// libEGL, which loads every EGL vendor, and Mesa's pulls in LLVM again.

fn driver_env(node: u64) -> Vec<(&'static str, std::ffi::OsString)> {

    let set = |vars: &[&str]| vars.iter().any(|v| std::env::var_os(v).is_some());

    let (major, minor) = (rustix::fs::major(node), rustix::fs::minor(node));

    let Ok(driver) = std::fs::read_link(format!("/sys/dev/char/{major}:{minor}/device/driver"))

    else {

        return Vec::new();

    };

    let driver = driver

        .file_name()

        .and_then(|d| d.to_str())

        .unwrap_or_default();

    let mut env = Vec::new();

    let glob = match driver {

        "nvidia" => Some("*nvidia*"),

        "amdgpu" | "radeon" => Some("*radeon*"),

        "i915" | "xe" => Some("*intel*"),

        "nouveau" => Some("*nouveau*"),

        _ => None,

    };

    let loader = [

        "VK_LOADER_DRIVERS_SELECT",

        "VK_LOADER_DRIVERS_DISABLE",

        "VK_DRIVER_FILES",

        "VK_ICD_FILENAMES",

        "VK_ADD_DRIVER_FILES",

    ];

    if let Some(glob) = glob

        && !set(&loader)

    {

        env.push(("VK_LOADER_DRIVERS_SELECT", glob.into()));

    }

    let glvnd = [

        "__EGL_VENDOR_LIBRARY_FILENAMES",

        "__EGL_VENDOR_LIBRARY_DIRS",

    ];

    if driver == "nvidia"

        && !set(&glvnd)

        && let Some(vendor) = nvidia_egl_vendor()

    {

        env.push(("__EGL_VENDOR_LIBRARY_FILENAMES", vendor.into()));

    }

    env

}

/// NVIDIA's glvnd EGL vendor file, from glvnd's default search path.

fn nvidia_egl_vendor() -> Option<std::path::PathBuf> {

    [

        "/run/opengl-driver/share/glvnd/egl_vendor.d",

        "/etc/glvnd/egl_vendor.d",

        "/usr/share/glvnd/egl_vendor.d",

    ]

    .into_iter()

    .filter_map(|dir| std::fs::read_dir(dir).ok())

    .flatten()

    .filter_map(Result::ok)

    .map(|e| e.path())

    .find(|p| {

        p.extension().is_some_and(|e| e == "json")

            && p.file_name()

                .and_then(|n| n.to_str())

                .is_some_and(|n| n.contains("nvidia"))

    })

}

/// The render node of `physical`, or its primary node when it has none.

unsafe fn drm_node(instance: &ash::Instance, physical: vk::PhysicalDevice) -> Option<u64> {

    let mut drm = vk::PhysicalDeviceDrmPropertiesEXT::default();

    let mut props = vk::PhysicalDeviceProperties2::default().push_next(&mut drm);

    unsafe { instance.get_physical_device_properties2(physical, &mut props) };

    let node = |major: i64, minor: i64| rustix::fs::makedev(major as u32, minor as u32);

    if drm.has_render == vk::TRUE {

        Some(node(drm.render_major, drm.render_minor))

    } else if drm.has_primary == vk::TRUE {

        Some(node(drm.primary_major, drm.primary_minor))

    } else {

        None

    }

}

unsafe fn pick_physical(

    instance: &ash::Instance,

    wanted: Option<&str>,

    node: Option<u64>,

) -> Option<(vk::PhysicalDevice, u32, String)> {

    let devices = unsafe { instance.enumerate_physical_devices() }.ok()?;

    let mut candidates = Vec::new();

    for physical in devices {

        let props = unsafe { instance.get_physical_device_properties(physical) };

        if props.api_version < vk::API_VERSION_1_3 {

            continue;

        }

        let name = props

            .device_name_as_c_str()

            .unwrap_or(c"unknown")

            .to_string_lossy()

            .into_owned();

        if wanted.is_some_and(|w| !name.contains(w)) {

            continue;

        }

        let mut f13 = vk::PhysicalDeviceVulkan13Features::default();

        let mut f2 = vk::PhysicalDeviceFeatures2::default().push_next(&mut f13);

        unsafe { instance.get_physical_device_features2(physical, &mut f2) };

        if f13.dynamic_rendering == vk::FALSE || f13.synchronization2 == vk::FALSE {

            continue;

        }

        let extensions =

            unsafe { instance.enumerate_device_extension_properties(physical) }.unwrap_or_default();

        let has = |n: &std::ffi::CStr| {

            extensions

                .iter()

                .any(|e| e.extension_name_as_c_str().ok() == Some(n))

        };

        if !has(ash::khr::push_descriptor::NAME) {

            continue;

        }

        if let Some(node) = node

            && !(has(ash::ext::physical_device_drm::NAME)

                && unsafe { serves_node(instance, physical, node) })

        {

            continue;

        }

        let families = unsafe { instance.get_physical_device_queue_family_properties(physical) };

        let Some(family) = families

            .iter()

            .position(|f| f.queue_flags.contains(vk::QueueFlags::GRAPHICS))

        else {

            continue;

        };

        let rank = match props.device_type {

            vk::PhysicalDeviceType::DISCRETE_GPU => 0,

            vk::PhysicalDeviceType::INTEGRATED_GPU => 1,

            vk::PhysicalDeviceType::VIRTUAL_GPU => 2,

            _ => 3,

        };

        candidates.push((rank, physical, family as u32, name));

    }

    candidates.sort_by_key(|c| c.0);

    candidates.into_iter().next().map(|(_, p, f, n)| (p, f, n))

}