wonfy-tools 0.2.0

mod score;

pub(crate) use score::{SearchParams, find_best_overlap};

use crate::error::StitchError;

pub(crate) struct GpuContext {
    pub(crate) device: wgpu::Device,
    pub(crate) queue: wgpu::Queue,
    pub(crate) preprocess_pipeline: wgpu::ComputePipeline,
    pub(crate) sad_pipeline: wgpu::ComputePipeline,
}

#[cfg(not(target_arch = "wasm32"))]
pub(crate) type SharedGpuContext = &'static GpuContext;
#[cfg(target_arch = "wasm32")]
pub(crate) type SharedGpuContext = std::rc::Rc<GpuContext>;

impl GpuContext {
    async fn new() -> Result<Self, StitchError> {
        let instance =
            wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle_from_env());

        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                ..Default::default()
            })
            .await
            .map_err(|err| StitchError::NoAdapter(err.to_string()))?;

        let (device, queue) = adapter
            .request_device(&wgpu::DeviceDescriptor {
                label: Some("wonfy-stitcher"),
                ..Default::default()
            })
            .await
            .map_err(|err| StitchError::Gpu(err.to_string()))?;

        let preprocess_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("stitcher-preprocess"),
            source: wgpu::ShaderSource::Wgsl(include_str!("preprocess.wgsl").into()),
        });

        let sad_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("stitcher-sad"),
            source: wgpu::ShaderSource::Wgsl(include_str!("sad.wgsl").into()),
        });

        let create_pipeline = |label, module| {
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some(label),
                layout: None,
                module,
                entry_point: Some("main"),
                compilation_options: Default::default(),
                cache: None,
            })
        };

        Ok(Self {
            preprocess_pipeline: create_pipeline("stitcher-preprocess", &preprocess_module),
            sad_pipeline: create_pipeline("stitcher-sad", &sad_module),
            device,
            queue,
        })
    }

    pub(crate) async fn read_buffer_u32(
        &self,
        buffer: &wgpu::Buffer,
    ) -> Result<Vec<u32>, StitchError> {
        let (sender, receiver) = futures_channel::oneshot::channel();

        buffer.map_async(wgpu::MapMode::Read, .., move |result| {
            let _ = sender.send(result);
        });

        #[cfg(not(target_arch = "wasm32"))]
        self.device
            .poll(wgpu::PollType::wait_indefinitely())
            .map_err(|err| StitchError::Gpu(format!("device poll failed: {err:?}")))?;

        receiver
            .await
            .map_err(|_| StitchError::Gpu("buffer map callback dropped".into()))?
            .map_err(|err| StitchError::Gpu(format!("buffer map failed: {err:?}")))?;

        let data = {
            let view = buffer
                .get_mapped_range(..)
                .map_err(|err| StitchError::Gpu(format!("buffer map range failed: {err:?}")))?;
            bytemuck::cast_slice::<u8, u32>(&view).to_vec()
        };
        buffer.unmap();

        Ok(data)
    }
}

#[cfg(not(target_arch = "wasm32"))]
pub(crate) async fn shared_context() -> Result<SharedGpuContext, StitchError> {
    static CONTEXT: std::sync::OnceLock<GpuContext> = std::sync::OnceLock::new();

    if let Some(context) = CONTEXT.get() {
        return Ok(context);
    }

    let context = GpuContext::new().await?;
    Ok(CONTEXT.get_or_init(|| context))
}

#[cfg(target_arch = "wasm32")]
pub(crate) async fn shared_context() -> Result<SharedGpuContext, StitchError> {
    use std::cell::RefCell;
    use std::rc::Rc;

    thread_local! {
        static CONTEXT: RefCell<Option<Rc<GpuContext>>> = const { RefCell::new(None) };
    }

    if let Some(context) = CONTEXT.with(|slot| slot.borrow().clone()) {
        return Ok(context);
    }

    let context = Rc::new(GpuContext::new().await?);
    CONTEXT.with(|slot| *slot.borrow_mut() = Some(context.clone()));
    Ok(context)
}