Moved in mod and rayon the specific code

Author: gents83

src/decoder.rs

@@ -8,7 +8,7 @@ use crate::parser::{
 };
 use crate::read_u8;
 use crate::upsampler::Upsampler;
-use crate::worker::{with_worker, PreferWorkerKind, RowData, Worker};
+use crate::worker::{compute_image_parallel, with_worker, PreferWorkerKind, RowData, Worker};
 use alloc::borrow::ToOwned;
 use alloc::sync::Arc;
 use alloc::vec::Vec;
@@ -1180,71 +1180,7 @@ fn compute_image(
     }
 }
 
-#[cfg(all(
-    not(any(target_arch = "asmjs", target_arch = "wasm32")),
-    feature = "rayon"
-))]
-fn compute_image_parallel(
-    components: &[Component],
-    data: Vec<Vec<u8>>,
-    output_size: Dimensions,
-    is_jfif: bool,
-    color_transform: Option<AdobeColorTransform>,
-) -> Result<Vec<u8>> {
-    use rayon::prelude::*;
-
-    let color_convert_func = choose_color_convert_func(components.len(), is_jfif, color_transform)?;
-    let upsampler = Upsampler::new(components, output_size.width, output_size.height)?;
-    let line_size = output_size.width as usize * components.len();
-    let mut image = vec![0u8; line_size * output_size.height as usize];
-
-    image
-        .par_chunks_mut(line_size)
-        .with_max_len(1)
-        .enumerate()
-        .for_each(|(row, line)| {
-            upsampler.upsample_and_interleave_row(
-                &data,
-                row,
-                output_size.width as usize,
-                line,
-                color_convert_func,
-            );
-        });
-
-    Ok(image)
-}
-
-#[cfg(any(
-    any(target_arch = "asmjs", target_arch = "wasm32"),
-    not(feature = "rayon")
-))]
-fn compute_image_parallel(
-    components: &[Component],
-    data: Vec<Vec<u8>>,
-    output_size: Dimensions,
-    is_jfif: bool,
-    color_transform: Option<AdobeColorTransform>,
-) -> Result<Vec<u8>> {
-    let color_convert_func = choose_color_convert_func(components.len(), is_jfif, color_transform)?;
-    let upsampler = Upsampler::new(components, output_size.width, output_size.height)?;
-    let line_size = output_size.width as usize * components.len();
-    let mut image = vec![0u8; line_size * output_size.height as usize];
-
-    for (row, line) in image.chunks_mut(line_size).enumerate() {
-        upsampler.upsample_and_interleave_row(
-            &data,
-            row,
-            output_size.width as usize,
-            line,
-            color_convert_func,
-        );
-    }
-
-    Ok(image)
-}
-
-fn choose_color_convert_func(
+pub(crate) fn choose_color_convert_func(
     component_count: usize,
     _is_jfif: bool,
     color_transform: Option<AdobeColorTransform>,

src/worker/mod.rs

@@ -6,8 +6,10 @@ mod multithreaded;
 ))]
 mod rayon;
 
+use crate::decoder::choose_color_convert_func;
 use crate::error::Result;
-use crate::parser::Component;
+use crate::parser::{AdobeColorTransform, Component, Dimensions};
+use crate::upsampler::Upsampler;
 use alloc::sync::Arc;
 use alloc::vec::Vec;
 
@@ -50,3 +52,38 @@ pub fn with_worker<T>(prefer: PreferWorkerKind, f: impl FnOnce(&mut dyn Worker)
         _ => self::immediate::with_immediate(f),
     }
 }
+
+pub fn compute_image_parallel(
+    components: &[Component],
+    data: Vec<Vec<u8>>,
+    output_size: Dimensions,
+    is_jfif: bool,
+    color_transform: Option<AdobeColorTransform>,
+) -> Result<Vec<u8>> {
+    #[cfg(all(
+        not(any(target_arch = "asmjs", target_arch = "wasm32")),
+        feature = "rayon"
+    ))]
+    return rayon::compute_image_parallel(components, data, output_size, is_jfif, color_transform);
+
+    #[allow(unreachable_code)]
+    {
+        let color_convert_func =
+            choose_color_convert_func(components.len(), is_jfif, color_transform)?;
+        let upsampler = Upsampler::new(components, output_size.width, output_size.height)?;
+        let line_size = output_size.width as usize * components.len();
+        let mut image = vec![0u8; line_size * output_size.height as usize];
+
+        for (row, line) in image.chunks_mut(line_size).enumerate() {
+            upsampler.upsample_and_interleave_row(
+                &data,
+                row,
+                output_size.width as usize,
+                line,
+                color_convert_func,
+            );
+        }
+
+        Ok(image)
+    }
+}

src/worker/rayon.rs

@@ -1,9 +1,14 @@
 use core::convert::TryInto;
 
-use crate::decoder::MAX_COMPONENTS;
+use rayon::iter::{IndexedParallelIterator, ParallelIterator};
+use rayon::slice::ParallelSliceMut;
+
+use crate::decoder::choose_color_convert_func;
 use crate::error::Result;
 use crate::idct::dequantize_and_idct_block;
-use crate::parser::Component;
+use crate::parser::{AdobeColorTransform, Component};
+use crate::upsampler::Upsampler;
+use crate::{decoder::MAX_COMPONENTS, parser::Dimensions};
 
 use std::sync::Arc;
 
@@ -166,8 +171,8 @@ impl super::Worker for Scoped {
                 let quantization_table = inner.quantization_tables[index].as_ref().unwrap().clone();
 
                 inner.offsets[index] += metadata.bytes_used();
-                let (result_block, tail) = core::mem::take(&mut result_blocks[index])
-                    .split_at_mut(metadata.bytes_used());
+                let (result_block, tail) =
+                    core::mem::take(&mut result_blocks[index]).split_at_mut(metadata.bytes_used());
                 result_blocks[index] = tail;
 
                 scope.spawn(move |_| {
@@ -190,3 +195,32 @@ impl ComponentMetadata {
         self.block_count * self.dct_scale * self.dct_scale
     }
 }
+
+pub fn compute_image_parallel(
+    components: &[Component],
+    data: Vec<Vec<u8>>,
+    output_size: Dimensions,
+    is_jfif: bool,
+    color_transform: Option<AdobeColorTransform>,
+) -> Result<Vec<u8>> {
+    let color_convert_func = choose_color_convert_func(components.len(), is_jfif, color_transform)?;
+    let upsampler = Upsampler::new(components, output_size.width, output_size.height)?;
+    let line_size = output_size.width as usize * components.len();
+    let mut image = vec![0u8; line_size * output_size.height as usize];
+
+    image
+        .par_chunks_mut(line_size)
+        .with_max_len(1)
+        .enumerate()
+        .for_each(|(row, line)| {
+            upsampler.upsample_and_interleave_row(
+                &data,
+                row,
+                output_size.width as usize,
+                line,
+                color_convert_func,
+            );
+        });
+
+    Ok(image)
+}