styx_codec/decoder/raw/
rgba.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 rgba_row_to_rgb24_neon(src: &[u8], dst: &mut [u8], width: usize) {
11    use std::arch::aarch64::{uint8x16x3_t, vld4q_u8, vst3q_u8};
12    debug_assert!(src.len() >= width * 4);
13    debug_assert!(dst.len() >= width * 3);
14
15    let src_ptr = src.as_ptr();
16    let dst_ptr = dst.as_mut_ptr();
17
18    let mut x = 0usize;
19    while x + 16 <= width {
20        unsafe {
21            let rgba = vld4q_u8(src_ptr.add(x * 4));
22            let rgb = uint8x16x3_t(rgba.0, rgba.1, rgba.2);
23            vst3q_u8(dst_ptr.add(x * 3), rgb);
24        }
25        x += 16;
26    }
27    for x in x..width {
28        unsafe {
29            let si = x * 4;
30            let di = x * 3;
31            *dst_ptr.add(di) = *src_ptr.add(si);
32            *dst_ptr.add(di + 1) = *src_ptr.add(si + 1);
33            *dst_ptr.add(di + 2) = *src_ptr.add(si + 2);
34        }
35    }
36}
37
38/// RGBA → RGB24 decoder (drops alpha).
39pub struct RgbaToRgbDecoder {
40    descriptor: CodecDescriptor,
41    pool: BufferPool,
42}
43
44impl RgbaToRgbDecoder {
45    pub fn new(max_width: u32, max_height: u32) -> Self {
46        let bytes = max_width as usize * max_height as usize * 3;
47        Self::with_input(
48            BufferPool::with_limits(2, bytes, 4),
49            FourCc::new(*b"RGBA"),
50            "rgba-strip",
51        )
52    }
53
54    pub fn with_pool(pool: BufferPool) -> Self {
55        Self::with_input(pool, FourCc::new(*b"RGBA"), "rgba-strip")
56    }
57
58    pub fn with_input(pool: BufferPool, input: FourCc, impl_name: &'static str) -> Self {
59        Self {
60            descriptor: CodecDescriptor {
61                kind: CodecKind::Decoder,
62                input,
63                output: FourCc::new(*b"RG24"),
64                name: "rgba2rgb",
65                impl_name,
66            },
67            pool,
68        }
69    }
70
71    pub fn with_input_for_max(
72        input: FourCc,
73        impl_name: &'static str,
74        max_width: u32,
75        max_height: u32,
76    ) -> Self {
77        let bytes = max_width as usize * max_height as usize * 3;
78        Self::with_input(BufferPool::with_limits(2, bytes, 4), input, impl_name)
79    }
80
81    /// Decode into a caller-provided tightly-packed RGB24 buffer.
82    ///
83    /// `dst` must be at least `width * height * 3` bytes.
84    pub fn decode_into(&self, input: &FrameLease, dst: &mut [u8]) -> Result<FrameMeta, CodecError> {
85        let meta = input.meta();
86        if meta.format.code != self.descriptor.input {
87            return Err(CodecError::FormatMismatch {
88                expected: self.descriptor.input,
89                actual: meta.format.code,
90            });
91        }
92        let plane = input
93            .planes()
94            .into_iter()
95            .next()
96            .ok_or_else(|| CodecError::Codec("rgba frame missing plane".into()))?;
97
98        let width = meta.format.resolution.width.get() as usize;
99        let height = meta.format.resolution.height.get() as usize;
100        let stride = plane.stride().max(width * 4);
101        let required = stride
102            .checked_mul(height)
103            .ok_or_else(|| CodecError::Codec("rgba stride overflow".into()))?;
104        if plane.data().len() < required {
105            return Err(CodecError::Codec("rgba plane buffer too short".into()));
106        }
107
108        let row_bytes = width * 3;
109        let out_len = row_bytes
110            .checked_mul(height)
111            .ok_or_else(|| CodecError::Codec("rgba output overflow".into()))?;
112        if dst.len() < out_len {
113            return Err(CodecError::Codec("rgba dst buffer too short".into()));
114        }
115
116        let src = plane.data();
117        dst[..out_len]
118            .par_chunks_mut(row_bytes)
119            .enumerate()
120            .for_each(|(y, dst_line)| {
121                let src_line = &src[y * stride..][..width * 4];
122                #[cfg(target_arch = "aarch64")]
123                unsafe {
124                    rgba_row_to_rgb24_neon(src_line, dst_line, width);
125                    return;
126                }
127                #[cfg(not(target_arch = "aarch64"))]
128                {
129                    for (dst_px, src_px) in dst_line
130                        .chunks_exact_mut(3)
131                        .zip(src_line.chunks_exact(4))
132                    {
133                        dst_px[0] = src_px[0];
134                        dst_px[1] = src_px[1];
135                        dst_px[2] = src_px[2];
136                    }
137                }
138            });
139
140        Ok(FrameMeta::new(
141            MediaFormat::new(
142                self.descriptor.output,
143                meta.format.resolution,
144                meta.format.color,
145            ),
146            meta.timestamp,
147        ))
148    }
149}
150
151impl Codec for RgbaToRgbDecoder {
152    fn descriptor(&self) -> &CodecDescriptor {
153        &self.descriptor
154    }
155
156    fn process(&self, input: FrameLease) -> Result<FrameLease, CodecError> {
157        let layout = plane_layout_from_dims(
158            input.meta().format.resolution.width,
159            input.meta().format.resolution.height,
160            3,
161        );
162        let mut buf = self.pool.lease();
163        unsafe { buf.resize_uninit(layout.len) };
164        let meta = self.decode_into(&input, buf.as_mut_slice())?;
165
166        Ok(unsafe {
167            FrameLease::single_plane_uninit(
168                meta,
169                buf,
170                layout.len,
171                layout.stride,
172            )
173        })
174    }
175}
176
177#[cfg(feature = "image")]
178impl ImageDecode for RgbaToRgbDecoder {
179    fn decode_image(&self, frame: FrameLease) -> Result<image::DynamicImage, CodecError> {
180        process_to_dynamic(self, frame)
181    }
182}