jpeg_decoder/worker/
mod.rs

1mod immediate;
2mod multithreaded;
3#[cfg(all(
4    not(any(target_arch = "asmjs", target_arch = "wasm32")),
5    feature = "rayon"
6))]
7mod rayon;
8
9use crate::decoder::{choose_color_convert_func, ColorTransform};
10use crate::error::Result;
11use crate::parser::{Component, Dimensions};
12use crate::upsampler::Upsampler;
13
14use alloc::sync::Arc;
15use alloc::vec::Vec;
16use core::cell::RefCell;
17
18pub struct RowData {
19    pub index: usize,
20    pub component: Component,
21    pub quantization_table: Arc<[u16; 64]>,
22}
23
24pub trait Worker {
25    fn start(&mut self, row_data: RowData) -> Result<()>;
26    fn append_row(&mut self, row: (usize, Vec<i16>)) -> Result<()>;
27    fn get_result(&mut self, index: usize) -> Result<Vec<u8>>;
28    /// Default implementation for spawning multiple tasks.
29    fn append_rows(&mut self, row: &mut dyn Iterator<Item = (usize, Vec<i16>)>) -> Result<()> {
30        for item in row {
31            self.append_row(item)?;
32        }
33        Ok(())
34    }
35}
36
37#[allow(dead_code)]
38pub enum PreferWorkerKind {
39    Immediate,
40    Multithreaded,
41}
42
43#[derive(Default)]
44pub struct WorkerScope {
45    inner: core::cell::RefCell<Option<WorkerScopeInner>>,
46}
47
48enum WorkerScopeInner {
49    #[cfg(all(
50        not(any(target_arch = "asmjs", target_arch = "wasm32")),
51        feature = "rayon"
52    ))]
53    Rayon(Box<rayon::Scoped>),
54    #[cfg(not(any(target_arch = "asmjs", target_arch = "wasm32")))]
55    Multithreaded(multithreaded::MpscWorker),
56    Immediate(immediate::ImmediateWorker),
57}
58
59impl WorkerScope {
60    pub fn with<T>(with: impl FnOnce(&Self) -> T) -> T {
61        with(&WorkerScope {
62            inner: RefCell::default(),
63        })
64    }
65
66    pub fn get_or_init_worker<T>(
67        &self,
68        prefer: PreferWorkerKind,
69        f: impl FnOnce(&mut dyn Worker) -> T,
70    ) -> T {
71        let mut inner = self.inner.borrow_mut();
72        let inner = inner.get_or_insert_with(move || match prefer {
73            #[cfg(all(
74                not(any(target_arch = "asmjs", target_arch = "wasm32")),
75                feature = "rayon"
76            ))]
77            PreferWorkerKind::Multithreaded => WorkerScopeInner::Rayon(Default::default()),
78            #[allow(unreachable_patterns)]
79            #[cfg(not(any(target_arch = "asmjs", target_arch = "wasm32")))]
80            PreferWorkerKind::Multithreaded => WorkerScopeInner::Multithreaded(Default::default()),
81            _ => WorkerScopeInner::Immediate(Default::default()),
82        });
83
84        f(match &mut *inner {
85            #[cfg(all(
86                not(any(target_arch = "asmjs", target_arch = "wasm32")),
87                feature = "rayon"
88            ))]
89            WorkerScopeInner::Rayon(worker) => worker.as_mut(),
90            #[cfg(not(any(target_arch = "asmjs", target_arch = "wasm32")))]
91            WorkerScopeInner::Multithreaded(worker) => worker,
92            WorkerScopeInner::Immediate(worker) => worker,
93        })
94    }
95}
96
97pub fn compute_image_parallel(
98    components: &[Component],
99    data: Vec<Vec<u8>>,
100    output_size: Dimensions,
101    color_transform: ColorTransform,
102) -> Result<Vec<u8>> {
103    #[cfg(all(
104        not(any(target_arch = "asmjs", target_arch = "wasm32")),
105        feature = "rayon"
106    ))]
107    return rayon::compute_image_parallel(components, data, output_size, color_transform);
108
109    #[allow(unreachable_code)]
110    {
111        let color_convert_func = choose_color_convert_func(components.len(), color_transform)?;
112        let upsampler = Upsampler::new(components, output_size.width, output_size.height)?;
113        let line_size = output_size.width as usize * components.len();
114        let mut image = vec![0u8; line_size * output_size.height as usize];
115
116        for (row, line) in image.chunks_mut(line_size).enumerate() {
117            upsampler.upsample_and_interleave_row(
118                &data,
119                row,
120                output_size.width as usize,
121                line,
122                color_convert_func,
123            );
124        }
125
126        Ok(image)
127    }
128}