styx_codec/decoder/raw/
mono.rs

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    // NEON loads u16 lanes; require row pointer to be 2-byte aligned.
45    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
74/// Monochrome 8-bit → RGB24 (channel replicate).
75pub 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    /// Decode into a caller-provided tightly-packed RGB24 buffer.
100    ///
101    /// `dst` must be at least `width * height * 3` bytes.
102    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
296/// Monochrome 16-bit little-endian → RGB24 (downshift + replicate).
297pub 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    /// Decode into a caller-provided tightly-packed RGB24 buffer.
322    ///
323    /// `dst` must be at least `width * height * 3` bytes.
324    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}