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: (usize, Vec) 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}
129