From 05f0d3d10dc883493782b070108a01ff89a3f719 Mon Sep 17 00:00:00 2001
From: bendn <bend.n@outlook.com>
Date: Sat, 27 Jan 2024 11:55:09 +0700
Subject: [PATCH] wgpu types

---
 Cargo.toml          |   4 +-
 src/lib.rs          |   2 +
 src/wgpu_convert.rs | 105 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 110 insertions(+), 1 deletion(-)
 create mode 100644 src/wgpu_convert.rs

diff --git a/Cargo.toml b/Cargo.toml
index 64af882..9222508 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "fimg"
-version = "0.4.34"
+version = "0.4.36"
 authors = ["bend-n <bend.n@outlook.com>"]
 license = "MIT"
 edition = "2021"
@@ -23,6 +23,7 @@ minifb = { version = "0.25.0", default-features = false, features = [
     "x11",
     "wayland",
 ], optional = true }
+wgpu = { version = "0.19.1", default-features = false, optional = true }
 
 [dev-dependencies]
 iai = { git = "https://github.com/bend-n/iai.git" }
@@ -59,6 +60,7 @@ text = ["fontdue"]
 blur = ["slur"]
 real-show = ["minifb", "text"]
 default = ["save", "scale"]
+wgpu-convert = ["dep:wgpu"]
 
 [profile.release]
 debug = 2
diff --git a/src/lib.rs b/src/lib.rs
index e6b060f..b829320 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -92,6 +92,8 @@ mod pack;
 mod span;
 mod sub;
 pub mod uninit;
+#[cfg(feature = "wgpu-convert")]
+mod wgpu_convert;
 pub use pack::Pack;
 pub mod pixels;
 #[cfg(feature = "scale")]
diff --git a/src/wgpu_convert.rs b/src/wgpu_convert.rs
new file mode 100644
index 0000000..c95346f
--- /dev/null
+++ b/src/wgpu_convert.rs
@@ -0,0 +1,105 @@
+use std::num::NonZeroU32;
+
+use wgpu::{util::*, *};
+
+use crate::Image;
+
+impl<T, const N: usize> Image<T, N> {
+    /// Get the size as a [`wgpu::Extend3d`].
+    pub fn wgpu_size(&self) -> Extent3d {
+        Extent3d {
+            width: self.width(),
+            height: self.height(),
+            depth_or_array_layers: 1,
+        }
+    }
+}
+
+impl<T: AsRef<[u8]>> Image<T, 4> {
+    /// Upload this image to the gpu, returning a [`wgpu::Texture`].
+    pub fn send(&self, dev: &Device, q: &Queue, usage: TextureUsages) -> Texture {
+        dev.create_texture_with_data(
+            &q,
+            &TextureDescriptor {
+                label: None,
+                size: self.wgpu_size(),
+                mip_level_count: 1,
+                sample_count: 1,
+                dimension: TextureDimension::D2,
+                format: TextureFormat::Rgba8Unorm,
+                view_formats: &[],
+                usage,
+            },
+            util::TextureDataOrder::LayerMajor,
+            self.bytes(),
+        )
+    }
+}
+
+impl Image<Box<[u8]>, 4> {
+    /// Downlodas a purportedly [`TextureFormat::Rgba8Unorm`] image from the gpu.
+    /// # Panics
+    ///
+    /// When a "error occurs while trying to async map a buffer".
+    pub fn download(
+        dev: &Device,
+        q: &Queue,
+        texture: &Texture,
+        (width, height): (NonZeroU32, NonZeroU32),
+    ) -> Self {
+        let mut encoder = dev.create_command_encoder(&CommandEncoderDescriptor { label: None });
+        let texture_size = Extent3d {
+            width: width.get(),
+            height: height.get(),
+            depth_or_array_layers: 1,
+        };
+
+        let row = width.get() as usize * 4;
+        let pad = {
+            let padding = (256 - row % 256) % 256;
+            row + padding
+        };
+
+        let output_buffer = dev.create_buffer(&BufferDescriptor {
+            label: None,
+            size: pad as u64 * height.get() as u64,
+            usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
+            mapped_at_creation: false,
+        });
+
+        encoder.copy_texture_to_buffer(
+            wgpu::ImageCopyTexture {
+                aspect: wgpu::TextureAspect::All,
+                texture,
+                mip_level: 0,
+                origin: wgpu::Origin3d::ZERO,
+            },
+            wgpu::ImageCopyBuffer {
+                buffer: &output_buffer,
+                layout: wgpu::ImageDataLayout {
+                    offset: 0,
+                    bytes_per_row: Some(pad as u32),
+                    rows_per_image: Some(height.get()),
+                },
+            },
+            texture_size,
+        );
+        q.submit(Some(encoder.finish()));
+
+        let buffer_slice = output_buffer.slice(..);
+        buffer_slice.map_async(wgpu::MapMode::Read, Result::unwrap);
+
+        dev.poll(wgpu::Maintain::Wait);
+
+        let mut out = crate::uninit::Image::<_, 4>::new(width, height);
+        for (padded, pixels) in buffer_slice
+            .get_mapped_range()
+            .chunks_exact(pad)
+            .zip(out.buf().chunks_exact_mut(row))
+        {
+            ::core::mem::MaybeUninit::write_slice(pixels, &padded[..row]);
+        }
+
+        unsafe { out.assume_init().boxed() }
+    }
+}