path: root/src/gpu.rs

use std::time::Duration;

use ash::vk;
use texture_share_vk_client::base::{
    vk_device::VkDevice, vk_entry::VkEntry, vk_instance::VkInstance, vk_setup::VkSetup,
};
use texture_share_vk_client::VkClient;

const TSV_SERVER_EXECUTABLE: &str = "/usr/bin/texture-share-vk-server";
const TSV_SERVER_SOCKET_PATH: &str = "/tmp/vk_server/vk_server.sock";
const TSV_SERVER_LOCK_PATH: &str = "/tmp/vk_server/vk_server.lock";
const TSV_SHMEM_PREFIX: &str = "shmem_img_";

/// Creates a wgpu instance configured for Vulkan.
pub fn create_instance() -> wgpu::Instance {
    wgpu::Instance::new(wgpu::InstanceDescriptor {
        backends: wgpu::Backends::VULKAN,
        flags: wgpu::InstanceFlags::default(),
        backend_options: wgpu::BackendOptions::default(),
        memory_budget_thresholds: wgpu::MemoryBudgetThresholds::default(),
        display: None,
    })
}

/// Requests a wgpu adapter, optionally compatible with a surface.
pub fn create_adapter(
    instance: &wgpu::Instance,
    surface: Option<&wgpu::Surface<'_>>,
) -> wgpu::Adapter {
    pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
        power_preference: wgpu::PowerPreference::HighPerformance,
        compatible_surface: surface,
        force_fallback_adapter: false,
    }))
    .expect("find vulkan adapter")
}

/// Requests a wgpu device + queue.
pub fn create_device(adapter: &wgpu::Adapter) -> (wgpu::Device, wgpu::Queue) {
    create_device_with_features(adapter, wgpu::Features::empty())
}

/// Requests a wgpu device + queue with additional features.
pub fn create_device_with_features(
    adapter: &wgpu::Adapter,
    features: wgpu::Features,
) -> (wgpu::Device, wgpu::Queue) {
    pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
        label: Some("device"),
        required_features: features,
        ..Default::default()
    }))
    .expect("create device")
}

/// Builds a texture-share-vk VkSetup by importing raw Vulkan handles from a wgpu device.
/// The returned VkSetup does NOT own the Vulkan objects (import_only = true).
fn import_vk_setup(device: &wgpu::Device) -> VkSetup {
    // Extract raw handles from the wgpu HAL device
    let (raw_instance, raw_device, raw_physical_device, raw_queue, queue_family_index, queue_index) = unsafe {
        let hal = device
            .as_hal::<wgpu_hal::api::Vulkan>()
            .expect("vulkan HAL device");

        let instance_shared = hal.shared_instance();
        let raw_instance = instance_shared.raw_instance().handle();
        let raw_device = hal.raw_device().handle();
        let raw_physical_device = hal.raw_physical_device();
        let raw_queue = hal.raw_queue();
        let queue_family_index = hal.queue_family_index();
        let queue_index = hal.queue_index();

        (
            raw_instance,
            raw_device,
            raw_physical_device,
            raw_queue,
            queue_family_index,
            queue_index,
        )
    };

    let vk_entry = Box::new(VkEntry::new().expect("vulkan entry"));
    let vk_instance =
        VkInstance::import_vk(Some(vk_entry), raw_instance, true).expect("import vulkan instance");
    let vk_device = VkDevice::import_vk(
        &vk_instance,
        raw_device,
        raw_physical_device,
        raw_queue,
        queue_family_index,
        queue_index,
        true,
    )
    .expect("import vulkan device");

    VkSetup::new(vk_instance, vk_device)
}

/// Creates a VkClient connected to the texture-share-vk server,
/// launching the server if needed.
pub fn create_tsv_client(device: &wgpu::Device) -> VkClient {
    let vk_setup = import_vk_setup(device);
    let timeout = Duration::from_secs(2);

    VkClient::new_with_server_launch(
        TSV_SERVER_SOCKET_PATH,
        Box::new(vk_setup),
        Duration::from_secs(1),
        TSV_SERVER_EXECUTABLE,
        TSV_SERVER_LOCK_PATH,
        TSV_SERVER_SOCKET_PATH,
        TSV_SHMEM_PREFIX,
        timeout,
        timeout,
        timeout,
        timeout,
        timeout,
    )
    .expect("connect to texture-share-vk server")
}

/// Re-export ImgFormat and ImgType for use in binaries.
pub use texture_share_vk_client::base::ipc::platform::img_data::{ImgFormat, ImgType};

/// Image format for texture-share-vk, matching RGBA8.
pub const TSV_FORMAT: ImgFormat = ImgFormat::R8G8B8A8;

/// Convert a TSV ImgFormat to a wgpu TextureFormat.
pub fn img_format_to_wgpu(fmt: ImgFormat) -> wgpu::TextureFormat {
    match fmt {
        ImgFormat::R8G8B8A8 => wgpu::TextureFormat::Rgba8Unorm,
        ImgFormat::R8G8B8 => wgpu::TextureFormat::Rgba8Unorm,
        ImgFormat::B8G8R8A8 => wgpu::TextureFormat::Bgra8Unorm,
        ImgFormat::B8G8R8 => wgpu::TextureFormat::Bgra8Unorm,
        ImgFormat::BC1_RGBA => wgpu::TextureFormat::Bc1RgbaUnorm,
        ImgFormat::BC3_RGBA => wgpu::TextureFormat::Bc3RgbaUnorm,
        ImgFormat::BC7_RGBA => wgpu::TextureFormat::Bc7RgbaUnorm,
        ImgFormat::Undefined => wgpu::TextureFormat::Rgba8Unorm,
    }
}

/// Returns the wgpu Features required for a TSV image format.
pub fn img_format_features(fmt: ImgFormat) -> wgpu::Features {
    if fmt.is_compressed() {
        wgpu::Features::TEXTURE_COMPRESSION_BC
    } else {
        wgpu::Features::empty()
    }
}

/// Extracts the raw vk::Image handle from a wgpu Texture.
///
/// # Safety
/// The returned handle must not outlive the texture.
pub unsafe fn raw_image(texture: &wgpu::Texture) -> vk::Image {
    let hal = texture
        .as_hal::<wgpu_hal::api::Vulkan>()
        .expect("vulkan HAL texture");
    hal.raw_handle()
}

/// Creates a vk::Fence on the VkClient's device.
pub fn create_fence(client: &VkClient) -> vk::Fence {
    let device = &client.get_vk_setup().device.device;
    let create_info = vk::FenceCreateInfo::default();
    unsafe { device.create_fence(&create_info, None) }.expect("create fence")
}

/// Destroys a vk::Fence.
pub fn destroy_fence(client: &VkClient, fence: vk::Fence) {
    let device = &client.get_vk_setup().device.device;
    unsafe { device.destroy_fence(fence, None) };
}

/// Refresh all texture inputs from TSV shared images.
/// Rebuilds the bind group if any texture was resized. Returns true if rebind occurred.
pub fn refresh_textures(
    cache: &mut crate::uniform::UniformCache,
    device: &wgpu::Device,
    client: &mut VkClient,
    fence: vk::Fence,
) -> bool {
    let mut needs_rebind = false;

    for texture in cache.textures_mut() {
        let tsv_name = match texture.tsv_name() {
            Some(n) => n.to_string(),
            None => continue,
        };

        match (
            texture.tsv_registered,
            client.find_image(&tsv_name, !texture.tsv_registered),
        ) {
            (_, Err(e)) => {
                log::error!("tsv find '{tsv_name}': {e}");
                continue;
            }
            (_, Ok(None)) => {
                log::warn!("tsv find '{tsv_name}': not found");
                continue;
            }
            (false, _) | (_, Ok(Some(true))) => match client.find_image_data(&tsv_name, true) {
                Ok(Some((_lock, data))) => {
                    let format = img_format_to_wgpu(data.format);
                    let view_dimension =
                        img_type_to_view_dimension(data.image_type, data.depth_or_array_layers);
                    if texture.resize(
                        device,
                        data.width,
                        data.height,
                        data.depth_or_array_layers,
                        view_dimension,
                        format,
                    ) {
                        needs_rebind = true;
                    }
                    texture.tsv_registered = true;
                }
                Ok(None) => {
                    log::warn!("recv_image '{tsv_name}': disappeared, re-registering");
                    texture.tsv_registered = false;
                    continue;
                }
                Err(e) => {
                    log::warn!("recv_image '{tsv_name}': {e}, re-registering");
                    texture.tsv_registered = false;
                    continue;
                }
            },
            _ => {}
        };

        let raw = unsafe { raw_image(texture.texture()) };
        match client.recv_image(
            &tsv_name,
            raw,
            vk::ImageLayout::UNDEFINED,
            vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL,
            fence,
        ) {
            Ok(Some(())) => {}
            Ok(None) => {
                log::warn!("recv_image '{tsv_name}': disappeared, re-registering");
                texture.tsv_registered = false;
            }
            Err(e) => {
                log::warn!("recv_image '{tsv_name}': {e}, re-registering");
                texture.tsv_registered = false;
            }
        }
    }

    if needs_rebind {
        cache.rebuild_bind_group(device);
    }

    needs_rebind
}

pub fn img_type_to_wgpu(image_type: ImgType, depth_or_array_layers: u32) -> wgpu::TextureDimension {
    match image_type {
        ImgType::D3 => wgpu::TextureDimension::D3,
        ImgType::D2 if depth_or_array_layers > 1 => wgpu::TextureDimension::D3,
        ImgType::D2 => wgpu::TextureDimension::D2,
    }
}

fn img_type_to_view_dimension(
    image_type: ImgType,
    depth_or_array_layers: u32,
) -> wgpu::TextureViewDimension {
    match image_type {
        ImgType::D3 => wgpu::TextureViewDimension::D3,
        ImgType::D2 if depth_or_array_layers > 1 => wgpu::TextureViewDimension::D2Array,
        _ => wgpu::TextureViewDimension::D2,
    }
}