ranga 1.0.0

//! GPU buffer wrapper for pixel data upload/download.
//!
//! [`GpuBuffer`] manages a wgpu storage buffer that holds pixel data on the GPU.
//! It provides upload and download helpers that work directly with [`PixelBuffer`].

use super::context::{GpuContext, GpuError};
use crate::pixel::{PixelBuffer, PixelFormat};

/// A GPU-side pixel buffer for compute operations.
///
/// Wraps a [`wgpu::Buffer`] with upload/download helpers for [`PixelBuffer`].
/// The buffer is created with storage, copy-src, and copy-dst usage flags
/// so it can be used in compute shaders and read back to the CPU.
///
/// # Examples
///
/// ```no_run
/// use ranga::gpu::{GpuContext, GpuBuffer};
/// use ranga::pixel::{PixelBuffer, PixelFormat};
///
/// let ctx = GpuContext::new().unwrap();
/// let buf = PixelBuffer::zeroed(256, 256, PixelFormat::Rgba8);
/// let gpu_buf = GpuBuffer::upload(&ctx, &buf);
/// let result = gpu_buf.download(&ctx).unwrap();
/// assert_eq!(result.width(), 256);
/// ```
pub struct GpuBuffer {
    buffer: wgpu::Buffer,
    size: u64,
    width: u32,
    height: u32,
    format: PixelFormat,
}

impl GpuBuffer {
    /// Upload a [`PixelBuffer`] to GPU storage.
    ///
    /// Creates a new GPU buffer initialized with the pixel data using
    /// [`mabda::buffer::create_storage_buffer`]. The buffer is usable as
    /// both a storage buffer (for compute shaders) and as a copy
    /// source/destination (for readback).
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use ranga::gpu::{GpuContext, GpuBuffer};
    /// use ranga::pixel::{PixelBuffer, PixelFormat};
    ///
    /// let ctx = GpuContext::new().unwrap();
    /// let buf = PixelBuffer::zeroed(64, 64, PixelFormat::Rgba8);
    /// let gpu_buf = GpuBuffer::upload(&ctx, &buf);
    /// assert_eq!(gpu_buf.width(), 64);
    /// ```
    #[must_use]
    pub fn upload(ctx: &GpuContext, buf: &PixelBuffer) -> Self {
        let buffer =
            mabda::buffer::create_storage_buffer(ctx.device(), &buf.data, "ranga_pixel_buf", false);
        Self {
            size: buf.data.len() as u64,
            width: buf.width,
            height: buf.height,
            format: buf.format,
            buffer,
        }
    }

    /// Download GPU buffer back to a [`PixelBuffer`].
    ///
    /// Uses [`mabda::buffer::read_buffer`] for synchronous GPU readback,
    /// then constructs a new [`PixelBuffer`] with the result.
    ///
    /// # Errors
    ///
    /// Returns [`GpuError::BufferOp`] if the buffer mapping or copy fails.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use ranga::gpu::{GpuContext, GpuBuffer};
    /// use ranga::pixel::{PixelBuffer, PixelFormat};
    ///
    /// let ctx = GpuContext::new().unwrap();
    /// let buf = PixelBuffer::zeroed(32, 32, PixelFormat::Rgba8);
    /// let gpu_buf = GpuBuffer::upload(&ctx, &buf);
    /// let result = gpu_buf.download(&ctx).unwrap();
    /// assert_eq!(result.data().len(), 32 * 32 * 4);
    /// ```
    #[must_use = "returns the downloaded pixel buffer"]
    pub fn download(&self, ctx: &GpuContext) -> Result<PixelBuffer, GpuError> {
        let data = mabda::buffer::read_buffer(ctx.device(), ctx.queue(), &self.buffer, self.size)
            .map_err(|e| GpuError::BufferOp(e.to_string()))?;

        PixelBuffer::new(data, self.width, self.height, self.format)
            .map_err(|e| GpuError::BufferOp(e.to_string()))
    }

    /// Start an asynchronous download of this GPU buffer.
    ///
    /// Returns immediately with a [`std::sync::mpsc::Receiver`] that will
    /// eventually contain the downloaded [`PixelBuffer`]. The caller can
    /// perform other work while the GPU readback is in progress, then call
    /// [`Receiver::recv`](std::sync::mpsc::Receiver::recv) to block until
    /// the result is ready.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use ranga::gpu::{GpuContext, GpuBuffer};
    /// use ranga::pixel::{PixelBuffer, PixelFormat};
    ///
    /// let ctx = GpuContext::new().unwrap();
    /// let buf = PixelBuffer::zeroed(64, 64, PixelFormat::Rgba8);
    /// let gpu_buf = GpuBuffer::upload(&ctx, &buf);
    /// let receiver = gpu_buf.download_async(&ctx);
    /// // ... do other work ...
    /// let result = receiver.recv().unwrap().unwrap();
    /// assert_eq!(result.width(), 64);
    /// ```
    #[must_use = "returns a receiver for the downloaded pixel buffer"]
    pub fn download_async(
        &self,
        ctx: &GpuContext,
    ) -> std::sync::mpsc::Receiver<Result<PixelBuffer, GpuError>> {
        let (tx, rx) = std::sync::mpsc::channel();
        let width = self.width;
        let height = self.height;
        let format = self.format;

        // mabda::PendingReadback::finish() requires &Device, which is not Send.
        // Use manual staging buffer approach so the map callback thread can
        // complete independently without needing the device.
        let staging = ctx.device().create_buffer(&wgpu::BufferDescriptor {
            label: Some("staging_readback_async"),
            size: self.size,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let mut encoder = ctx
            .device()
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        encoder.copy_buffer_to_buffer(&self.buffer, 0, &staging, 0, self.size);
        ctx.queue().submit(Some(encoder.finish()));

        let size = self.size as usize;
        let slice = staging.slice(..);
        let (map_tx, map_rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |result| {
            let _ = map_tx.send(result);
        });

        std::thread::spawn(move || {
            let map_result = map_rx.recv();
            let result = (|| -> Result<PixelBuffer, GpuError> {
                map_result
                    .map_err(|e| GpuError::BufferOp(e.to_string()))?
                    .map_err(|e| GpuError::BufferOp(e.to_string()))?;
                let data = {
                    let view = staging.slice(..);
                    let mapped = view.get_mapped_range();
                    mapped[..size].to_vec()
                };
                drop(staging);
                PixelBuffer::new(data, width, height, format)
                    .map_err(|e| GpuError::BufferOp(e.to_string()))
            })();
            let _ = tx.send(result);
        });

        rx
    }

    /// Access the underlying [`wgpu::Buffer`].
    #[must_use]
    #[inline]
    pub fn wgpu_buffer(&self) -> &wgpu::Buffer {
        &self.buffer
    }

    /// Size of the buffer in bytes.
    #[must_use]
    #[inline]
    pub fn size(&self) -> u64 {
        self.size
    }

    /// Width of the pixel buffer in pixels.
    #[must_use]
    #[inline]
    pub fn width(&self) -> u32 {
        self.width
    }

    /// Height of the pixel buffer in pixels.
    #[must_use]
    #[inline]
    pub fn height(&self) -> u32 {
        self.height
    }

    /// Pixel format of the buffer.
    #[must_use]
    #[inline]
    pub fn format(&self) -> PixelFormat {
        self.format
    }
}