Polyfill rayon in WASM

Author: kornelski

Cargo.toml

@@ -19,15 +19,17 @@ name = "dssim"
 path = "src/main.rs"
 
 [dependencies]
-dssim-core = { path = "./dssim-core", version = "3.1" }
+dssim-core = { path = "./dssim-core", version = "3.2", default-features = false }
 imgref = "1.9.1"
 getopts = "0.2.21"
-rayon = "1.5.1"
+rayon = { version = "1.5.1", optional = true }
 rgb = "0.8.31"
 lodepng = "3.5.2"
 load_image = { version = "2.16.1", features = ["lcms2-static"] }
 
 [features]
+default = ["threads", "dssim-core/default"]
+threads = ["rayon"]
 avif = ["load_image/avif"]
 webp = ["load_image/webp"]
 webp-static = ["load_image/webp-static"]

README.md

@@ -88,3 +88,11 @@ DSSIM is dual-licensed under [AGPL](LICENSE) or [commercial](https://supso.org/p
 * The lightness component of SSIM is ignored when comparing color channels.
 * SSIM score is pooled using mean absolute deviation. You can get per-pixel SSIM from the API to implement custom pooling.
 
+
+## Compiling for WASM
+
+For compatibility with single-threaded WASM runtimes, disable the `threads` Cargo feature. It's enabled by default, so to disable it, disable default features:
+
+```toml
+dssim-core = { version = "3.2", default-features = false }
+```

dssim-core/Cargo.toml

@@ -17,13 +17,17 @@ edition = "2018"
 crate-type = ["lib", "staticlib"]
 
 [dependencies]
-rayon = "1.5.0"
 imgref = "1.9.1"
 itertools = "0.10.3"
+rayon = { version = "1.5.1", optional = true }
 rgb = "0.8.31"
 
 [dev-dependencies]
 lodepng = "3.6.0"
 
+[features]
+default = ["threads"]
+threads = ["rayon"]
+
 [package.metadata.docs.rs]
 targets = ["x86_64-unknown-linux-gnu"]

dssim-core/src/dssim.rs

@@ -26,6 +26,8 @@ pub use crate::tolab::ToLABBitmap;
 pub use crate::val::Dssim as Val;
 use imgref::*;
 use itertools::multizip;
+#[cfg(not(feature = "threads"))]
+use crate::lieon as rayon;
 use rayon::prelude::*;
 use rgb::{RGB, RGBA};
 use std::borrow::Borrow;
@@ -293,8 +295,8 @@ impl Dssim {
     /// `Val` is a fancy wrapper for `f64`
     pub fn compare<M: Borrow<DssimImage<f32>>>(&self, original_image: &DssimImage<f32>, modified_image: M) -> (Val, Vec<SsimMap>) {
         let modified_image = modified_image.borrow();
-        let res: Vec<_> = self.scale_weights.par_iter().cloned().zip(
-            modified_image.scale.par_iter().zip(original_image.scale.par_iter())
+        let res: Vec<_> = self.scale_weights.as_slice().par_iter().cloned().zip(
+            modified_image.scale.as_slice().par_iter().zip(original_image.scale.as_slice().par_iter())
         ).enumerate().map(|(n, (weight, (modified_image_scale, original_image_scale)))| {
             let scale_width = original_image_scale.chan[0].width;
             let scale_height = original_image_scale.chan[0].height;
@@ -395,9 +397,9 @@ impl Dssim {
         debug_assert_eq!(img1_img2_blur.len(), original.mu.len());
         debug_assert_eq!(img1_img2_blur.len(), original.img_sq_blur.len());
 
-        let mu_iter = original.mu.par_iter().cloned().zip_eq(modified.mu.par_iter().cloned());
-        let sq_iter = original.img_sq_blur.par_iter().cloned().zip_eq(modified.img_sq_blur.par_iter().cloned());
-        img1_img2_blur.par_iter().cloned().zip_eq(mu_iter).zip_eq(sq_iter).zip_eq(map_out.par_iter_mut())
+        let mu_iter = original.mu.as_slice().par_iter().cloned().zip_eq(modified.mu.as_slice().par_iter().cloned());
+        let sq_iter = original.img_sq_blur.as_slice().par_iter().cloned().zip_eq(modified.img_sq_blur.as_slice().par_iter().cloned());
+        img1_img2_blur.par_iter().cloned().zip_eq(mu_iter).zip_eq(sq_iter).zip_eq(map_out.as_mut_slice().par_iter_mut())
         .for_each(|(((img1_img2_blur, (mu1, mu2)), (img1_sq_blur, img2_sq_blur)), map_out)| {
             let mu1mu1 = mu1 * mu1;
             let mu1mu2 = mu1 * mu2;

dssim-core/src/lib.rs

@@ -14,6 +14,8 @@ mod image;
 mod linear;
 mod tolab;
 mod val;
+#[cfg(not(feature = "threads"))]
+mod lieon;
 
 pub use crate::dssim::*;
 pub use crate::image::*;

dssim-core/src/lieon.rs

@@ -0,0 +1,76 @@
+//! Shim for single-threaded rayon replacement
+
+pub mod prelude {
+    pub use super::*;
+}
+
+pub trait ParSliceLie<T> {
+    fn par_chunks(&self, n: usize) -> std::slice::Chunks<'_, T>;
+}
+
+pub trait ParSliceMutLie<T> {
+    fn par_chunks_mut(&mut self, n: usize) -> std::slice::ChunksMut<'_, T>;
+}
+
+pub trait ParIntoIterLie<T> {
+    type IntoIter;
+    fn into_par_iter(self) -> Self::IntoIter;
+}
+
+pub trait ParIterLie<T> {
+    type Iter;
+    fn par_iter(&self) -> Self::Iter;
+}
+
+pub trait ParIterMutLie<'a, T> {
+    type Iter;
+    fn par_iter_mut(&'a mut self) -> Self::Iter;
+}
+
+pub fn join<A, B>(a: impl FnOnce() -> A, b: impl FnOnce() -> B) -> (A, B) {
+    let a = a();
+    let b = b();
+    (a, b)
+}
+
+impl<'a, T> ParSliceLie<T> for &'a [T] {
+    fn par_chunks(&self, n: usize) -> std::slice::Chunks<'_, T> {
+        self.chunks(n)
+    }
+}
+
+impl<'a, T> ParSliceLie<T> for &'a mut [T] {
+    fn par_chunks(&self, n: usize) -> std::slice::Chunks<'_, T> {
+        self.chunks(n)
+    }
+}
+
+impl<'a, T> ParSliceMutLie<T> for &'a mut [T] {
+    fn par_chunks_mut(&mut self, n: usize) -> std::slice::ChunksMut<'_, T> {
+        self.chunks_mut(n)
+    }
+}
+
+impl<'a, T> ParIterLie<T> for &'a [T] {
+    type Iter = std::slice::Iter<'a, T>;
+
+    fn par_iter(&self) -> Self::Iter {
+        self.iter()
+    }
+}
+
+impl<'a, T> ParIterMutLie<'a, T> for &'a mut [T] {
+    type Iter = std::slice::IterMut<'a, T>;
+
+    fn par_iter_mut(&'a mut self) -> Self::Iter {
+        self.iter_mut()
+    }
+}
+
+
+impl<T> ParIntoIterLie<T> for Vec<T> {
+    type IntoIter = std::vec::IntoIter<T>;
+    fn into_par_iter(self) -> Self::IntoIter {
+        self.into_iter()
+    }
+}

dssim-core/src/tolab.rs

@@ -5,6 +5,8 @@ use crate::image::ToRGB;
 use crate::image::RGBAPLU;
 use crate::image::RGBLU;
 use imgref::*;
+#[cfg(not(feature = "threads"))]
+use crate::lieon as rayon;
 use rayon::prelude::*;
 
 const D65x: f64 = 0.9505;
@@ -69,7 +71,7 @@ impl ToLABBitmap for GBitmap {
         unsafe { out.set_len(size) };
 
         // For output width == stride
-        out.par_chunks_mut(width).enumerate().for_each(|(y, out_row)| {
+        out.as_mut_slice().par_chunks_mut(width).enumerate().for_each(|(y, out_row)| {
             let start = y * self.stride();
             let in_row = &self.buf()[start..start + width];
             let out_row = &mut out_row[0..width];
@@ -102,8 +104,8 @@ fn rgb_to_lab<T: Copy + Sync + Send + 'static, F>(img: ImgRef<'_, T>, cb: F) ->
     unsafe { out_b.set_len(size) };
 
     // For output width == stride
-    out_l.par_chunks_mut(width).zip(
-        out_a.par_chunks_mut(width).zip(out_b.par_chunks_mut(width))
+    out_l.as_mut_slice().par_chunks_mut(width).zip(
+        out_a.as_mut_slice().par_chunks_mut(width).zip(out_b.as_mut_slice().par_chunks_mut(width))
     ).enumerate()
     .for_each(|(y, (l_row, (a_row, b_row)))| {
         let start = y * stride;

src/main.rs

@@ -18,6 +18,7 @@
  */
 #![allow(clippy::manual_range_contains)]
 use getopts::Options;
+#[cfg(feature = "threads")]
 use rayon::prelude::*;
 use std::env;
 use std::path::PathBuf;
@@ -70,7 +71,12 @@ fn main() {
 
     let mut attr = dssim::Dssim::new();
 
-    let files = files.par_iter().map(|file| -> Result<_, String> {
+    #[cfg(feature = "threads")]
+    let files_iter = files.par_iter();
+    #[cfg(not(feature = "threads"))]
+    let files_iter = files.iter();
+
+    let files = files_iter.map(|file| -> Result<_, String> {
         let image = dssim::load_image(&attr, &file).map_err(|e| format!("Can't load {}, because: {}", file.display(), e))?;
         Ok((file, image))
     }).collect::<Result<Vec<_>,_>>();
@@ -102,7 +108,12 @@ fn main() {
         println!("{:.8}\t{}", dssim, file2.display());
 
         if map_output_file.is_some() {
-            ssim_maps.par_iter().enumerate().for_each(|(n, map_meta)| {
+            #[cfg(feature = "threads")]
+            let ssim_maps_iter = ssim_maps.par_iter();
+            #[cfg(not(feature = "threads"))]
+            let ssim_maps_iter = ssim_maps.iter();
+
+            ssim_maps_iter.enumerate().for_each(|(n, map_meta)| {
                 let avgssim = map_meta.ssim as f32;
                 let out: Vec<_> = map_meta.map.pixels().map(|ssim|{
                     let max = 1_f32 - ssim;