xgpui-gpu-interop 0.2.0

//! Adapter for wgpu engines running on Metal with their own device and queue.
use super::{SyncPoint, surface_texture};
use crate::{EngineSurface, SurfaceDesc, SurfaceFormat, SurfaceSlot};
use anyhow::{Context as _, Result};
use objc2::rc::Retained;
use objc2::runtime::ProtocolObject;
use objc2_metal::{MTLDevice, MTLSharedEvent, MTLTextureType, MTLTextureUsage};
use wgpu::hal::api::Metal;

/// Wraps surface slots as `wgpu::Texture`s and presents after submission.
///
/// `present` stages a shared-event signal on the engine queue, submits, then
/// publishes the slot, so the contract "present only after the signal is
/// submitted" holds by construction. Staged signals are queue-wide: do not
/// submit on this queue from other threads concurrently with `present`.
pub struct WgpuSurfaceTarget {
    surface: EngineSurface,
    device: wgpu::Device,
    queue: wgpu::Queue,
    event: Retained<ProtocolObject<dyn MTLSharedEvent>>,
    next: u64,
    textures: Vec<Option<wgpu::Texture>>,
    gpu_wait: bool,
    heavy: Option<crate::heavy::Heavy>,
    heavy_sync: Option<SyncPoint>,
}

// SAFETY: wgpu handles are Send + Sync on native targets; MTLSharedEvent is
// thread-safe for retain/release and is only staged on the engine's queue. The
// target is used from one engine thread at a time (`&mut self` methods).
unsafe impl Send for WgpuSurfaceTarget {}

impl WgpuSurfaceTarget {
    /// Features the engine's wgpu device must enable.
    pub fn required_features() -> wgpu::Features {
        wgpu::Features::empty()
    }

    /// Create a surface for `desc` and a target that renders into it.
    pub fn new(desc: SurfaceDesc, device: &wgpu::Device, queue: &wgpu::Queue) -> Result<Self> {
        let surface = EngineSurface::new(desc)?;
        let event = unsafe { device.as_hal::<Metal>() }
            .context("wgpu engine surfaces require the Metal backend")?
            .raw_device()
            .newSharedEvent()
            .context("MTLSharedEvent allocation failed")?;
        Ok(Self {
            surface,
            device: device.clone(),
            queue: queue.clone(),
            event,
            next: 0,
            textures: vec![None; desc.slots],
            gpu_wait: true,
            heavy: None,
            heavy_sync: None,
        })
    }

    pub fn format(&self) -> wgpu::TextureFormat {
        match self.surface.desc().format {
            SurfaceFormat::Bgra8Srgb => wgpu::TextureFormat::Bgra8Unorm,
            SurfaceFormat::Rgba16FloatLinearP3 => wgpu::TextureFormat::Rgba16Float,
        }
    }

    /// The engine's texture over `slot` (cached per slot).
    pub fn texture(&mut self, slot: &SurfaceSlot) -> Result<wgpu::Texture> {
        if let Some(texture) = &self.textures[slot.index()] {
            return Ok(texture.clone());
        }
        let desc = self.surface.desc();
        let raw_device = unsafe { self.device.as_hal::<Metal>() }
            .context("Metal hal device")?
            .raw_device()
            .clone();
        let raw = surface_texture(
            &raw_device,
            slot.iosurface(),
            &desc,
            MTLTextureUsage::RenderTarget
                | MTLTextureUsage::ShaderRead
                | MTLTextureUsage::ShaderWrite,
        )?;
        let size = wgpu::Extent3d {
            width: desc.width,
            height: desc.height,
            depth_or_array_layers: 1,
        };
        let hal = unsafe {
            wgpu::hal::metal::Device::texture_from_raw(
                raw,
                self.format(),
                MTLTextureType::Type2D,
                1,
                1,
                wgpu::hal::CopyExtent {
                    width: desc.width,
                    height: desc.height,
                    depth: 1,
                },
                None,
            )
        };
        let mut usage = wgpu::TextureUsages::RENDER_ATTACHMENT
            | wgpu::TextureUsages::TEXTURE_BINDING
            | wgpu::TextureUsages::COPY_SRC
            | wgpu::TextureUsages::COPY_DST;
        if desc.format.supports_storage() {
            usage |= wgpu::TextureUsages::STORAGE_BINDING;
        }
        let texture = unsafe {
            self.device.create_texture_from_hal::<Metal>(
                hal,
                &wgpu::TextureDescriptor {
                    label: Some("engine_surface_slot"),
                    size,
                    mip_level_count: 1,
                    sample_count: 1,
                    dimension: wgpu::TextureDimension::D2,
                    format: self.format(),
                    usage,
                    view_formats: &[],
                },
                wgpu::TextureUses::UNINITIALIZED,
            )
        };
        self.textures[slot.index()] = Some(texture.clone());
        Ok(texture)
    }

    /// Submit `commands`, signal completion, and publish `slot`.
    pub fn present(
        &mut self,
        slot: SurfaceSlot,
        commands: impl IntoIterator<Item = wgpu::CommandBuffer>,
    ) -> Result<()> {
        let submitted = std::time::Instant::now();
        self.next += 1;
        unsafe { self.queue.as_hal::<Metal>() }
            .context("Metal hal queue")?
            .add_signal_event(self.event.clone(), self.next);
        self.queue.submit(commands);
        let sync = if self.gpu_wait {
            // SAFETY: the signal was staged on the submission above.
            unsafe { SyncPoint::submitted(self.event.clone(), self.next) }
        } else {
            SyncPoint::new(self.event.clone(), self.next)
        };
        self.publish(slot, sync, submitted)
    }

    /// Whether the compositor may wait on the GPU for this target's frames
    /// (default). With `false` (observe-only) a frame is shown only once its
    /// GPU work is observed complete, so a slow engine frame never holds up
    /// composition; the last complete frame stays on screen meanwhile.
    /// Costs up to a frame of latency. Applies to frames presented after the
    /// call.
    pub fn set_gpu_wait(&mut self, enabled: bool) {
        self.gpu_wait = enabled;
    }

    /// See [`Self::set_gpu_wait`].
    pub fn gpu_wait(&self) -> bool {
        self.gpu_wait
    }

    /// Frames are always handed over with GPU synchronization on macOS.
    pub fn gpu_sync(&self) -> bool {
        true
    }

    /// Heavy-engine mode (see the `heavy` module): observe-only
    /// presentation, pacing on compositor lag ([`Self::pace`]) and chunk
    /// advice ([`Self::suggested_chunks`]). `None` restores the defaults
    /// (GPU-side wait, no pacing).
    pub fn set_heavy(&mut self, config: Option<crate::HeavyEngine>) {
        self.set_gpu_wait(config.is_none());
        self.heavy = config.map(crate::heavy::Heavy::new);
        self.heavy_sync = None;
    }

    /// Heavy-engine state, when enabled.
    pub fn heavy(&self) -> Option<crate::HeavyStatus> {
        self.heavy.as_ref().map(|heavy| heavy.status())
    }

    /// In heavy-engine mode, how long to wait before rendering the next
    /// frame (backs off while the compositor lags); zero otherwise. Also
    /// measures the previous frame's GPU time for [`Self::suggested_chunks`].
    pub fn pace(&mut self) -> std::time::Duration {
        let Some(heavy) = &mut self.heavy else {
            return std::time::Duration::ZERO;
        };
        if let Some(sync) = &self.heavy_sync {
            let reached = sync.reached();
            heavy.observe(reached);
            if reached {
                self.heavy_sync = None;
            }
        }
        heavy.pace(&self.surface.stats())
    }

    /// In heavy-engine mode, how many separate submissions to split a
    /// frame's work into so each stays near the chunk budget; 1 otherwise.
    /// Submit all but the last chunk with the engine's queue, the last with
    /// [`Self::present`].
    pub fn suggested_chunks(&self) -> u32 {
        self.heavy
            .as_ref()
            .map_or(1, |heavy| heavy.suggested_chunks())
    }

    fn publish(
        &mut self,
        slot: SurfaceSlot,
        sync: SyncPoint,
        submitted: std::time::Instant,
    ) -> Result<()> {
        if let Some(heavy) = &mut self.heavy {
            heavy.on_present(submitted);
            if self.heavy_sync.is_none() {
                self.heavy_sync = Some(sync.clone());
            }
        }
        self.surface.present(slot, sync)
    }

    pub fn surface(&self) -> &EngineSurface {
        &self.surface
    }
}