oxiui-compute-wgpu 0.1.0

//! Headless GPU compute context: `Instance` → `Adapter` → `Device` + `Queue`.
//!
//! [`ComputeContext`] performs the full no-window, no-surface initialisation
//! chain required for pure GPU compute workloads (sparse solvers, LBM, MC/DC,
//! …).  Three constructors plus a fluent [`ContextBuilder`] are provided:
//!
//! * [`ComputeContext::try_new`] — returns `Option<Self>`; `None` means no GPU
//!   adapter is available (graceful CI skip, never panics).
//! * [`ComputeContext::new`] — returns `Result<Self, ComputeError>`; exposes the
//!   underlying failure reason through [`ComputeError`].
//! * [`ComputeContext::new_async`] — async variant; awaits adapter and device
//!   requests directly without a `pollster::block_on` wrapper.
//! * [`ComputeContext::builder`] — returns a [`ContextBuilder`] for fluent
//!   configuration of limits, features, and power preference.
//!
//! Both sync constructors use `PowerPreference::HighPerformance` and
//! `wgpu::Limits::default()` (not `downlevel_defaults()`, which caps the
//! compute feature set).

use crate::error::ComputeError;

// ── ComputeContext ─────────────────────────────────────────────────────────────

/// An initialised headless GPU compute context.
///
/// Owns the logical [`wgpu::Device`], the associated [`wgpu::Queue`], and
/// the [`wgpu::AdapterInfo`] snapshot captured at construction time.  No
/// window handle, surface, or swap-chain is involved.
pub struct ComputeContext {
    /// The logical GPU device.
    pub device: wgpu::Device,
    /// The command submission queue for the device.
    pub queue: wgpu::Queue,
    /// Adapter metadata snapshot (vendor, backend, driver, …).
    adapter_info: wgpu::AdapterInfo,
}

impl ComputeContext {
    /// Return a reference to the adapter metadata captured at construction time.
    ///
    /// The returned [`wgpu::AdapterInfo`] contains fields such as `name`,
    /// `vendor`, `device`, `backend`, `driver`, and `driver_info`.
    ///
    /// ```rust,no_run
    /// use oxiui_compute_wgpu::ComputeContext;
    ///
    /// if let Some(ctx) = ComputeContext::try_new() {
    ///     let info = ctx.adapter_info();
    ///     println!("GPU backend: {:?}", info.backend);
    /// }
    /// ```
    pub fn adapter_info(&self) -> &wgpu::AdapterInfo {
        &self.adapter_info
    }

    /// Return a [`ContextBuilder`] for fluent configuration of limits,
    /// features, and power preference.
    ///
    /// ```rust,no_run
    /// use oxiui_compute_wgpu::ComputeContext;
    ///
    /// let ctx = ComputeContext::builder()
    ///     .with_power_preference(wgpu::PowerPreference::LowPower)
    ///     .build();
    /// ```
    pub fn builder() -> ContextBuilder {
        ContextBuilder::default()
    }

    /// Create a context with high-performance power preference and default limits.
    ///
    /// # Errors
    ///
    /// * [`ComputeError::NoAdapter`] — no suitable GPU adapter was found.
    /// * [`ComputeError::DeviceRequest`] — the device/queue request failed.
    ///
    /// ```rust,no_run
    /// use oxiui_compute_wgpu::{ComputeContext, ComputeError};
    ///
    /// match ComputeContext::new() {
    ///     Ok(ctx)                      => { let _ = ctx; }
    ///     Err(ComputeError::NoAdapter) => { /* skip */ }
    ///     Err(e)                       => panic!("unexpected: {e}"),
    /// }
    /// ```
    pub fn new() -> Result<Self, ComputeError> {
        ContextBuilder::default().build()
    }

    /// Try to create a `ComputeContext`, returning `None` when no suitable GPU
    /// adapter is available on this host.
    ///
    /// This constructor never panics.  Call sites that want a graceful skip on
    /// headless CI environments (VMs, containers without GPU pass-through) should
    /// use this variant:
    ///
    /// ```rust,no_run
    /// use oxiui_compute_wgpu::ComputeContext;
    ///
    /// if let Some(ctx) = ComputeContext::try_new() {
    ///     // GPU is available — run the compute workload
    ///     let _ = ctx;
    /// } else {
    ///     // No GPU — skip gracefully
    /// }
    /// ```
    pub fn try_new() -> Option<Self> {
        Self::new().ok()
    }

    /// Async variant of [`new`][Self::new] — awaits adapter and device requests
    /// directly without a `pollster::block_on` wrapper.
    ///
    /// Suitable for use inside an async runtime (Tokio, async-std, etc.).
    ///
    /// # Errors
    ///
    /// Same as [`new`][Self::new].
    ///
    /// ```rust,no_run
    /// use oxiui_compute_wgpu::ComputeContext;
    ///
    /// # async fn run() -> Result<(), oxiui_compute_wgpu::ComputeError> {
    /// let ctx = ComputeContext::new_async().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn new_async() -> Result<Self, ComputeError> {
        ContextBuilder::default().build_async().await
    }

    // ── Convenience delegates to ContextBuilder ─────────────────────────────

    /// Start building a context with custom memory limits.
    ///
    /// Equivalent to `ComputeContext::builder().with_limits(limits)`.
    pub fn with_limits(limits: wgpu::Limits) -> ContextBuilder {
        ContextBuilder::default().with_limits(limits)
    }

    /// Start building a context with specific GPU features enabled.
    ///
    /// Equivalent to `ComputeContext::builder().with_features(features)`.
    pub fn with_features(features: wgpu::Features) -> ContextBuilder {
        ContextBuilder::default().with_features(features)
    }

    /// Start building a context with a specific power preference.
    ///
    /// Equivalent to `ComputeContext::builder().with_power_preference(pref)`.
    pub fn with_power_preference(pref: wgpu::PowerPreference) -> ContextBuilder {
        ContextBuilder::default().with_power_preference(pref)
    }
}

// ── ContextBuilder ─────────────────────────────────────────────────────────────

/// Fluent builder for [`ComputeContext`].
///
/// Compose limits, features, and power preference in one chain, then call
/// [`build`][ContextBuilder::build] (sync) or [`build_async`][ContextBuilder::build_async]
/// (async) to finalise.
///
/// ```rust,no_run
/// use oxiui_compute_wgpu::{ComputeContext, ComputeError};
///
/// let result = ComputeContext::builder()
///     .with_power_preference(wgpu::PowerPreference::HighPerformance)
///     .with_limits(wgpu::Limits::default())
///     .build();
/// ```
#[derive(Debug, Default)]
pub struct ContextBuilder {
    power_preference: wgpu::PowerPreference,
    required_features: wgpu::Features,
    required_limits: Option<wgpu::Limits>,
}

impl ContextBuilder {
    /// Set the GPU power preference.
    ///
    /// Defaults to [`wgpu::PowerPreference::HighPerformance`] when not
    /// called.
    pub fn with_power_preference(mut self, pref: wgpu::PowerPreference) -> Self {
        self.power_preference = pref;
        self
    }

    /// Request optional GPU features (e.g. `TIMESTAMP_QUERY`, `SHADER_F16`).
    ///
    /// If the adapter does not support the requested features, [`build`] will
    /// return [`ComputeError::DeviceRequest`] with a descriptive message before
    /// attempting `request_device`.
    pub fn with_features(mut self, features: wgpu::Features) -> Self {
        self.required_features = features;
        self
    }

    /// Override the default device limits.
    ///
    /// Use [`wgpu::Limits::downlevel_defaults()`] for maximum compatibility or
    /// supply custom limits for high-throughput compute workloads.
    pub fn with_limits(mut self, limits: wgpu::Limits) -> Self {
        self.required_limits = Some(limits);
        self
    }

    /// Blocking variant: run the full adapter + device init on the current thread.
    ///
    /// # Errors
    ///
    /// * [`ComputeError::NoAdapter`] — no GPU adapter matched the options.
    /// * [`ComputeError::DeviceRequest`] — features or device request failed.
    pub fn build(self) -> Result<ComputeContext, ComputeError> {
        pollster::block_on(self.build_async())
    }

    /// Async variant: await adapter and device requests inside the caller's runtime.
    ///
    /// # Errors
    ///
    /// * [`ComputeError::NoAdapter`] — no GPU adapter matched the options.
    /// * [`ComputeError::DeviceRequest`] — features or device request failed.
    pub async fn build_async(self) -> Result<ComputeContext, ComputeError> {
        let instance = wgpu::Instance::default();

        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: self.power_preference,
                force_fallback_adapter: false,
                // No surface — pure compute, no swap-chain required.
                compatible_surface: None,
            })
            .await
            .map_err(|_| ComputeError::NoAdapter)?;

        // Pre-check requested features before attempting device acquisition so
        // callers get a clear error instead of a cryptic RequestDeviceError.
        if !self.required_features.is_empty()
            && !adapter.features().contains(self.required_features)
        {
            return Err(ComputeError::DeviceRequest(format!(
                "adapter does not support requested features: {:?}",
                self.required_features
            )));
        }

        // Capture adapter metadata before consuming the adapter.
        let adapter_info = adapter.get_info();

        let limits = self.required_limits.unwrap_or_default();

        let (device, queue) = adapter
            .request_device(&wgpu::DeviceDescriptor {
                label: Some("oxiui-compute-wgpu"),
                required_features: self.required_features,
                required_limits: limits,
                ..Default::default()
            })
            .await
            .map_err(|e| ComputeError::DeviceRequest(e.to_string()))?;

        Ok(ComputeContext {
            device,
            queue,
            adapter_info,
        })
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    // ── existing tests (preserved) ───────────────────────────────────────────

    #[test]
    fn try_new_does_not_panic() {
        // Gracefully skips if no GPU adapter — must never panic.
        let _ = ComputeContext::try_new();
    }

    #[test]
    fn new_returns_result() {
        match ComputeContext::new() {
            Ok(_ctx) => { /* GPU available — context created successfully */ }
            Err(ComputeError::NoAdapter) => {
                // No GPU on this host (CI, headless VM) — acceptable skip.
            }
            Err(ComputeError::DeviceRequest(ref msg)) => {
                panic!("unexpected DeviceRequest error: {msg}")
            }
            Err(e) => {
                panic!("unexpected error: {e}")
            }
        }
    }

    #[test]
    fn try_new_consistent_with_new() {
        // try_new() must be consistent with new(): both fail or both succeed.
        let via_new = ComputeContext::new();
        let via_try = ComputeContext::try_new();
        match (via_new, via_try) {
            (Ok(_), Some(_)) | (Err(_), None) => { /* consistent */ }
            (Ok(_), None) => panic!("new() succeeded but try_new() returned None"),
            (Err(e), Some(_)) => panic!("new() failed but try_new() returned Some: {e}"),
        }
    }

    // ── new tests (S1) ───────────────────────────────────────────────────────

    /// Non-GPU test: verify that `ContextBuilder::default()` constructs without
    /// panicking, even before `build()` is called.
    #[test]
    fn builder_chain_defaults() {
        // The builder itself must be constructable regardless of GPU availability.
        let _builder = ContextBuilder::default()
            .with_power_preference(wgpu::PowerPreference::HighPerformance)
            .with_limits(wgpu::Limits::default())
            .with_features(wgpu::Features::empty());
        // Attempt build; whether it succeeds depends on host GPU availability —
        // either outcome is acceptable.
        let _result = _builder.build();
        // No assertion: success or NoAdapter are both valid outcomes.
    }

    /// GPU-gated: adapter_info() returns a non-empty backend string.
    #[test]
    fn context_has_adapter_info() {
        oxiui_core::require_gpu!(ctx, ComputeContext::try_new());
        let info = ctx.adapter_info();
        let backend_str = format!("{:?}", info.backend);
        assert!(!backend_str.is_empty(), "backend string must not be empty");
    }

    /// GPU-gated: builder with LowPower preference builds successfully.
    #[test]
    fn builder_with_low_power() {
        oxiui_core::require_gpu!(
            ctx,
            ComputeContext::with_power_preference(wgpu::PowerPreference::LowPower)
                .build()
                .ok()
        );
        let _ = ctx;
    }

    /// GPU-gated: new_async() via pollster::block_on produces a valid context.
    #[test]
    fn new_async_via_pollster() {
        oxiui_core::require_gpu!(ctx, pollster::block_on(ComputeContext::new_async()).ok());
        let _ = ctx;
    }

    /// GPU-gated: requesting all features should return a clean error (not panic)
    /// when the adapter does not support them all.
    #[test]
    fn with_unsupported_features_returns_error() {
        // `wgpu::Features::all()` is almost certainly not fully supported on any
        // single adapter; we expect either a DeviceRequest error or a successful
        // build (the latter is allowed on hardware that does support everything).
        // What must NOT happen is a panic.
        let result = ComputeContext::with_features(wgpu::Features::all()).build();
        match result {
            Ok(_) => { /* hardware supports all features — acceptable */ }
            Err(ComputeError::NoAdapter) => { /* no GPU — skip */ }
            Err(ComputeError::DeviceRequest(_)) => { /* expected clean error */ }
            Err(e) => panic!("unexpected error variant: {e}"),
        }
    }
}