1use styx_core::prelude::*;
2
3#[cfg(feature = "image")]
4use crate::decoder::{ImageDecode, process_to_dynamic};
5use crate::{Codec, CodecDescriptor, CodecError, CodecKind};
6use rayon::prelude::*;
7
8#[cfg(target_arch = "aarch64")]
9#[inline(always)]
10unsafe fn mono8_row_to_rgb24_neon(src: &[u8], dst: &mut [u8], width: usize) {
11 use std::arch::aarch64::{uint8x16x3_t, vld1q_u8, vst3q_u8};
12 debug_assert!(src.len() >= width);
13 debug_assert!(dst.len() >= width * 3);
14
15 let src_ptr = src.as_ptr();
16 let dst_ptr = dst.as_mut_ptr();
17 let mut x = 0usize;
18 while x + 16 <= width {
19 unsafe {
20 let g = vld1q_u8(src_ptr.add(x));
21 let rgb = uint8x16x3_t(g, g, g);
22 vst3q_u8(dst_ptr.add(x * 3), rgb);
23 }
24 x += 16;
25 }
26 for x in x..width {
27 unsafe {
28 let gray = *src_ptr.add(x);
29 let o = x * 3;
30 *dst_ptr.add(o) = gray;
31 *dst_ptr.add(o + 1) = gray;
32 *dst_ptr.add(o + 2) = gray;
33 }
34 }
35}
36
37#[cfg(target_arch = "aarch64")]
38#[inline(always)]
39unsafe fn mono16_row_to_rgb24_neon(src: &[u8], dst: &mut [u8], width: usize) -> bool {
40 use std::arch::aarch64::{uint8x8x3_t, vld1q_u16, vshrn_n_u16, vst3_u8};
41 debug_assert!(src.len() >= width * 2);
42 debug_assert!(dst.len() >= width * 3);
43
44 if (src.as_ptr() as usize) & 0x1 != 0 {
46 return false;
47 }
48
49 let src_u16 = src.as_ptr() as *const u16;
50 let dst_ptr = dst.as_mut_ptr();
51 let mut x = 0usize;
52 while x + 8 <= width {
53 unsafe {
54 let v16 = vld1q_u16(src_u16.add(x));
55 let g8 = vshrn_n_u16(v16, 8);
56 vst3_u8(dst_ptr.add(x * 3), uint8x8x3_t(g8, g8, g8));
57 }
58 x += 8;
59 }
60 for x in x..width {
61 unsafe {
62 let si = x * 2;
63 let v = u16::from_le_bytes([*src.get_unchecked(si), *src.get_unchecked(si + 1)]);
64 let g = (v >> 8) as u8;
65 let di = x * 3;
66 *dst_ptr.add(di) = g;
67 *dst_ptr.add(di + 1) = g;
68 *dst_ptr.add(di + 2) = g;
69 }
70 }
71 true
72}
73
74pub struct Mono8ToRgbDecoder {
76 descriptor: CodecDescriptor,
77 pool: BufferPool,
78}
79
80impl Mono8ToRgbDecoder {
81 pub fn new(max_width: u32, max_height: u32) -> Self {
82 let bytes = max_width as usize * max_height as usize * 3;
83 Self::with_pool(BufferPool::with_limits(2, bytes, 4))
84 }
85
86 pub fn with_pool(pool: BufferPool) -> Self {
87 Self {
88 descriptor: CodecDescriptor {
89 kind: CodecKind::Decoder,
90 input: FourCc::new(*b"R8 "),
91 output: FourCc::new(*b"RG24"),
92 name: "mono2rgb",
93 impl_name: "mono8-replicate",
94 },
95 pool,
96 }
97 }
98
99 pub fn decode_into(&self, input: &FrameLease, dst: &mut [u8]) -> Result<FrameMeta, CodecError> {
103 let meta = input.meta();
104 if meta.format.code != self.descriptor.input {
105 return Err(CodecError::FormatMismatch {
106 expected: self.descriptor.input,
107 actual: meta.format.code,
108 });
109 }
110 let plane = input
111 .planes()
112 .into_iter()
113 .next()
114 .ok_or_else(|| CodecError::Codec("mono frame missing plane".into()))?;
115
116 let width = meta.format.resolution.width.get() as usize;
117 let height = meta.format.resolution.height.get() as usize;
118 let stride = plane.stride().max(width);
119 let required = stride
120 .checked_mul(height)
121 .ok_or_else(|| CodecError::Codec("mono stride overflow".into()))?;
122 if plane.data().len() < required {
123 return Err(CodecError::Codec("mono plane buffer too short".into()));
124 }
125
126 let row_bytes = width * 3;
127 let out_len = row_bytes
128 .checked_mul(height)
129 .ok_or_else(|| CodecError::Codec("mono output overflow".into()))?;
130 if dst.len() < out_len {
131 return Err(CodecError::Codec("mono dst buffer too short".into()));
132 }
133
134 let src = plane.data();
135 dst[..out_len]
136 .par_chunks_mut(row_bytes)
137 .enumerate()
138 .for_each(|(y, dst_line)| {
139 let src_line = &src[y * stride..][..width];
140 #[cfg(target_arch = "aarch64")]
141 unsafe {
142 mono8_row_to_rgb24_neon(src_line, dst_line, width);
143 return;
144 }
145 #[cfg(not(target_arch = "aarch64"))]
146 {
147 for (dst_px, &gray) in dst_line.chunks_exact_mut(3).zip(src_line.iter()) {
148 dst_px[0] = gray;
149 dst_px[1] = gray;
150 dst_px[2] = gray;
151 }
152 }
153 });
154
155 Ok(FrameMeta::new(
156 MediaFormat::new(
157 self.descriptor.output,
158 meta.format.resolution,
159 meta.format.color,
160 ),
161 meta.timestamp,
162 ))
163 }
164}
165
166impl Codec for Mono8ToRgbDecoder {
167 fn descriptor(&self) -> &CodecDescriptor {
168 &self.descriptor
169 }
170
171 fn process(&self, input: FrameLease) -> Result<FrameLease, CodecError> {
172 let layout = plane_layout_from_dims(
173 input.meta().format.resolution.width,
174 input.meta().format.resolution.height,
175 3,
176 );
177 let mut buf = self.pool.lease();
178 unsafe { buf.resize_uninit(layout.len) };
179 let meta = self.decode_into(&input, buf.as_mut_slice())?;
180
181 Ok(unsafe {
182 FrameLease::single_plane_uninit(
183 meta,
184 buf,
185 layout.len,
186 layout.stride,
187 )
188 })
189 }
190}
191
192#[cfg(feature = "image")]
193impl ImageDecode for Mono8ToRgbDecoder {
194 fn decode_image(&self, frame: FrameLease) -> Result<image::DynamicImage, CodecError> {
195 process_to_dynamic(self, frame)
196 }
197}
198
199#[cfg(test)]
200#[allow(clippy::items_after_test_module)]
201mod tests {
202 use super::*;
203
204 #[test]
205 fn mono8_expands_to_rgb24() {
206 let decoder = Mono8ToRgbDecoder::new(16, 4);
207 let res = Resolution::new(5, 3).unwrap();
208 let height = res.height.get() as usize;
209 let width = res.width.get() as usize;
210 let stride = 8usize;
211 let mut src = vec![0u8; stride * height];
212 for y in 0..height {
213 for x in 0..width {
214 src[y * stride + x] = (y as u8) * 10 + (x as u8);
215 }
216 }
217
218 let pool = BufferPool::with_limits(1, stride * height, 4);
219 let mut buf = pool.lease();
220 buf.resize(src.len());
221 buf.as_mut_slice().copy_from_slice(&src);
222
223 let frame = FrameLease::single_plane(
224 FrameMeta::new(
225 MediaFormat::new(FourCc::new(*b"R8 "), res, ColorSpace::Unknown),
226 0,
227 ),
228 buf,
229 stride * height,
230 stride,
231 );
232 let out = decoder.process(frame).unwrap();
233 assert_eq!(out.meta().format.code, FourCc::new(*b"RG24"));
234 let planes = out.planes();
235 assert_eq!(planes.len(), 1);
236 let data = planes[0].data();
237 assert_eq!(data.len(), width * height * 3);
238 for y in 0..height {
239 for x in 0..width {
240 let g = (y as u8) * 10 + (x as u8);
241 let o = (y * width + x) * 3;
242 assert_eq!(&data[o..o + 3], &[g, g, g]);
243 }
244 }
245 }
246
247 #[test]
248 fn mono16_expands_to_rgb24() {
249 let decoder = Mono16ToRgbDecoder::new(16, 4);
250 let res = Resolution::new(5, 3).unwrap();
251 let height = res.height.get() as usize;
252 let width = res.width.get() as usize;
253 let stride = 12usize;
254
255 let mut src = vec![0u8; stride * height];
256 for y in 0..height {
257 for x in 0..width {
258 let v = (((y * width + x) * 257) as u16).to_le_bytes();
259 let o = y * stride + x * 2;
260 src[o] = v[0];
261 src[o + 1] = v[1];
262 }
263 }
264
265 let pool = BufferPool::with_limits(1, src.len(), 4);
266 let mut buf = pool.lease();
267 buf.resize(src.len());
268 buf.as_mut_slice().copy_from_slice(&src);
269
270 let frame = FrameLease::single_plane(
271 FrameMeta::new(
272 MediaFormat::new(FourCc::new(*b"R16 "), res, ColorSpace::Unknown),
273 0,
274 ),
275 buf,
276 src.len(),
277 stride,
278 );
279 let out = decoder.process(frame).unwrap();
280 assert_eq!(out.meta().format.code, FourCc::new(*b"RG24"));
281 let planes = out.planes();
282 assert_eq!(planes.len(), 1);
283 let data = planes[0].data();
284 assert_eq!(data.len(), width * height * 3);
285 for y in 0..height {
286 for x in 0..width {
287 let v = ((y * width + x) * 257) as u16;
288 let g = (v >> 8) as u8;
289 let o = (y * width + x) * 3;
290 assert_eq!(&data[o..o + 3], &[g, g, g]);
291 }
292 }
293 }
294}
295
296pub struct Mono16ToRgbDecoder {
298 descriptor: CodecDescriptor,
299 pool: BufferPool,
300}
301
302impl Mono16ToRgbDecoder {
303 pub fn new(max_width: u32, max_height: u32) -> Self {
304 let bytes = max_width as usize * max_height as usize * 3;
305 Self::with_pool(BufferPool::with_limits(2, bytes, 4))
306 }
307
308 pub fn with_pool(pool: BufferPool) -> Self {
309 Self {
310 descriptor: CodecDescriptor {
311 kind: CodecKind::Decoder,
312 input: FourCc::new(*b"R16 "),
313 output: FourCc::new(*b"RG24"),
314 name: "mono2rgb",
315 impl_name: "mono16-replicate",
316 },
317 pool,
318 }
319 }
320
321 pub fn decode_into(&self, input: &FrameLease, dst: &mut [u8]) -> Result<FrameMeta, CodecError> {
325 let meta = input.meta();
326 if meta.format.code != self.descriptor.input {
327 return Err(CodecError::FormatMismatch {
328 expected: self.descriptor.input,
329 actual: meta.format.code,
330 });
331 }
332 let plane = input
333 .planes()
334 .into_iter()
335 .next()
336 .ok_or_else(|| CodecError::Codec("mono16 frame missing plane".into()))?;
337
338 let width = meta.format.resolution.width.get() as usize;
339 let height = meta.format.resolution.height.get() as usize;
340 let stride = plane.stride().max(width * 2);
341 let required = stride
342 .checked_mul(height)
343 .ok_or_else(|| CodecError::Codec("mono16 stride overflow".into()))?;
344 if plane.data().len() < required {
345 return Err(CodecError::Codec("mono16 plane buffer too short".into()));
346 }
347
348 let row_bytes = width * 3;
349 let out_len = row_bytes
350 .checked_mul(height)
351 .ok_or_else(|| CodecError::Codec("mono16 output overflow".into()))?;
352 if dst.len() < out_len {
353 return Err(CodecError::Codec("mono16 dst buffer too short".into()));
354 }
355
356 let src = plane.data();
357 dst[..out_len]
358 .par_chunks_mut(row_bytes)
359 .enumerate()
360 .for_each(|(y, dst_line)| {
361 let src_line = &src[y * stride..][..width * 2];
362 #[cfg(target_arch = "aarch64")]
363 unsafe {
364 if mono16_row_to_rgb24_neon(src_line, dst_line, width) {
365 return;
366 }
367 }
368 for (dst_px, chunk) in dst_line.chunks_exact_mut(3).zip(src_line.chunks_exact(2)) {
369 let gray16 = u16::from_le_bytes([chunk[0], chunk[1]]);
370 let gray8 = (gray16 >> 8) as u8;
371 dst_px[0] = gray8;
372 dst_px[1] = gray8;
373 dst_px[2] = gray8;
374 }
375 });
376
377 Ok(FrameMeta::new(
378 MediaFormat::new(
379 self.descriptor.output,
380 meta.format.resolution,
381 meta.format.color,
382 ),
383 meta.timestamp,
384 ))
385 }
386}
387
388impl Codec for Mono16ToRgbDecoder {
389 fn descriptor(&self) -> &CodecDescriptor {
390 &self.descriptor
391 }
392
393 fn process(&self, input: FrameLease) -> Result<FrameLease, CodecError> {
394 let layout = plane_layout_from_dims(
395 input.meta().format.resolution.width,
396 input.meta().format.resolution.height,
397 3,
398 );
399 let mut buf = self.pool.lease();
400 unsafe { buf.resize_uninit(layout.len) };
401 let meta = self.decode_into(&input, buf.as_mut_slice())?;
402
403 Ok(unsafe {
404 FrameLease::single_plane_uninit(
405 meta,
406 buf,
407 layout.len,
408 layout.stride,
409 )
410 })
411 }
412}
413
414#[cfg(feature = "image")]
415impl ImageDecode for Mono16ToRgbDecoder {
416 fn decode_image(&self, frame: FrameLease) -> Result<image::DynamicImage, CodecError> {
417 process_to_dynamic(self, frame)
418 }
419}