styx_codec/decoder/raw/
rgba.rs1use 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
38pub 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 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}