ezk_image/
planes.rs

1use crate::{
2    PixelFormat, StrictApi,
3    plane_decs::{
4        I01X_PLANES, I21X_PLANES, I41X_PLANES, I420_PLANES, I422_PLANES, I444_PLANES, NV12_PLANES,
5        P01X_PLANES, PlaneDesc,
6    },
7    util::ArrayIter,
8};
9use std::mem::MaybeUninit;
10
11#[derive(Debug, thiserror::Error)]
12#[error("got invalid number of planes, expected {expected} but only got {got}")]
13pub struct InvalidNumberOfPlanesError {
14    pub expected: usize,
15    pub got: usize,
16}
17
18pub(crate) fn read_planes<'a, const N: usize>(
19    mut iter: impl Iterator<Item = (&'a [u8], usize)>,
20) -> Result<[(&'a [u8], usize); N], InvalidNumberOfPlanesError> {
21    let mut out: [(&'a [u8], usize); N] = [(&[], 0); N];
22
23    for (i, out) in out.iter_mut().enumerate() {
24        *out = iter.next().ok_or(InvalidNumberOfPlanesError {
25            expected: N,
26            got: i,
27        })?;
28    }
29
30    Ok(out)
31}
32
33pub(crate) fn read_planes_mut<'a, const N: usize>(
34    mut iter: impl Iterator<Item = (&'a mut [u8], usize)>,
35) -> Result<[(&'a mut [u8], usize); N], InvalidNumberOfPlanesError> {
36    let mut out: [MaybeUninit<(&'a mut [u8], usize)>; N] = [const { MaybeUninit::uninit() }; N];
37
38    for (i, out) in out.iter_mut().enumerate() {
39        out.write(iter.next().ok_or(InvalidNumberOfPlanesError {
40            expected: N,
41            got: i,
42        })?);
43    }
44
45    Ok(out.map(|plane| unsafe { plane.assume_init() }))
46}
47
48/// Infer the planes for an image in the given format using the given dimensions and strides
49///
50/// # Panics
51///
52/// If `buf` is too small for the given dimensions this function will panic
53#[deny(clippy::arithmetic_side_effects)]
54pub fn infer<S: AnySlice>(
55    format: PixelFormat,
56    buf: S,
57    width: usize,
58    height: usize,
59    strides: Option<&[usize]>,
60) -> impl Iterator<Item = S> {
61    match format {
62        #[cfg(feature = "I420")]
63        PixelFormat::I420 => ArrayIter::from(infer_i420(buf, width, height, strides)),
64        #[cfg(feature = "I422")]
65        PixelFormat::I422 => ArrayIter::from(infer_i422(buf, width, height, strides)),
66        #[cfg(feature = "I444")]
67        PixelFormat::I444 => ArrayIter::from(infer_i444(buf, width, height, strides)),
68        #[cfg(feature = "I010")]
69        PixelFormat::I010 => ArrayIter::from(infer_i01x(buf, width, height, strides)),
70        #[cfg(feature = "I012")]
71        PixelFormat::I012 => ArrayIter::from(infer_i01x(buf, width, height, strides)),
72        #[cfg(feature = "I210")]
73        PixelFormat::I210 => ArrayIter::from(infer_i21x(buf, width, height, strides)),
74        #[cfg(feature = "I212")]
75        PixelFormat::I212 => ArrayIter::from(infer_i21x(buf, width, height, strides)),
76        #[cfg(feature = "I410")]
77        PixelFormat::I410 => ArrayIter::from(infer_i41x(buf, width, height, strides)),
78        #[cfg(feature = "I412")]
79        PixelFormat::I412 => ArrayIter::from(infer_i41x(buf, width, height, strides)),
80        #[cfg(feature = "NV12")]
81        PixelFormat::NV12 => ArrayIter::from(infer_nv12(buf, width, height, strides)),
82        #[cfg(feature = "P010")]
83        PixelFormat::P010 => ArrayIter::from(infer_p01x(buf, width, height, strides)),
84        #[cfg(feature = "P012")]
85        PixelFormat::P012 => ArrayIter::from(infer_p01x(buf, width, height, strides)),
86        #[cfg(feature = "YUYV")]
87        PixelFormat::YUYV => ArrayIter::from([buf]),
88        #[cfg(feature = "RGBA")]
89        PixelFormat::RGBA => ArrayIter::from([buf]),
90        #[cfg(feature = "BGRA")]
91        PixelFormat::BGRA => ArrayIter::from([buf]),
92        #[cfg(feature = "ARGB")]
93        PixelFormat::ARGB => ArrayIter::from([buf]),
94        #[cfg(feature = "ABGR")]
95        PixelFormat::ABGR => ArrayIter::from([buf]),
96        #[cfg(feature = "RGB")]
97        PixelFormat::RGB => ArrayIter::from([buf]),
98        #[cfg(feature = "BGR")]
99        PixelFormat::BGR => ArrayIter::from([buf]),
100    }
101}
102
103#[deny(clippy::arithmetic_side_effects)]
104fn infer_impl<const N: usize, S: AnySlice>(
105    plane_decs: [PlaneDesc; N],
106    mut buf: S,
107    width: usize,
108    height: usize,
109    strides: Option<&[usize]>,
110) -> [S; N] {
111    let strides = strides.map(|strides| <[usize; N]>::try_from(strides).unwrap());
112
113    // Infer default strides for a packed buffer
114    let strides: [usize; N] =
115        strides.unwrap_or_else(|| plane_decs.map(|desc| desc.packed_stride(width)));
116
117    let mut out: [MaybeUninit<S>; N] = [const { MaybeUninit::uninit() }; N];
118
119    for ((desc, stride), out) in plane_decs.into_iter().zip(strides).zip(out.iter_mut()) {
120        let split_at = desc.height_op.op(height).strict_mul_(stride);
121
122        let (prev, rem) = buf.slice_split_at(split_at);
123
124        out.write(prev);
125        buf = rem;
126    }
127
128    out.map(|p| unsafe { p.assume_init() })
129}
130
131/// Infer the planes for a full I420 image using the given dimensions
132///
133/// # Panics
134///
135/// If `buf` is too small for the given dimensions this function will panic
136#[deny(clippy::arithmetic_side_effects)]
137pub fn infer_i420<S: AnySlice>(
138    buf: S,
139    width: usize,
140    height: usize,
141    strides: Option<&[usize]>,
142) -> [S; 3] {
143    infer_impl(I420_PLANES, buf, width, height, strides)
144}
145
146/// Infer the planes for a full I422 image using the given dimensions
147///
148/// # Panics
149///
150/// If `buf` is too small for the given dimensions this function will panic
151#[deny(clippy::arithmetic_side_effects)]
152pub fn infer_i422<S: AnySlice>(
153    buf: S,
154    width: usize,
155    height: usize,
156    strides: Option<&[usize]>,
157) -> [S; 3] {
158    infer_impl(I422_PLANES, buf, width, height, strides)
159}
160
161/// Infer the planes for a full I444 image using the given dimensions
162///
163/// # Panics
164///
165/// If `buf` is too small for the given dimensions this function will panic
166#[deny(clippy::arithmetic_side_effects)]
167pub fn infer_i444<S: AnySlice>(
168    buf: S,
169    width: usize,
170    height: usize,
171    strides: Option<&[usize]>,
172) -> [S; 3] {
173    infer_impl(I444_PLANES, buf, width, height, strides)
174}
175
176/// Infer the planes for a full I010 or I012 image using the given dimensions
177///
178/// # Panics
179///
180/// If `buf` is too small for the given dimensions this function will panic
181#[deny(clippy::arithmetic_side_effects)]
182pub fn infer_i01x<S: AnySlice>(
183    buf: S,
184    width: usize,
185    height: usize,
186    strides: Option<&[usize]>,
187) -> [S; 3] {
188    infer_impl(I01X_PLANES, buf, width, height, strides)
189}
190
191/// Infer the planes for a full I210 or I212 image using the given dimensions
192///
193/// # Panics
194///
195/// If `buf` is too small for the given dimensions this function will panic
196#[deny(clippy::arithmetic_side_effects)]
197pub fn infer_i21x<S: AnySlice>(
198    buf: S,
199    width: usize,
200    height: usize,
201    strides: Option<&[usize]>,
202) -> [S; 3] {
203    infer_impl(I21X_PLANES, buf, width, height, strides)
204}
205
206/// Infer the planes for a full I410 or I412 image using the given dimensions
207///
208/// # Panics
209///
210/// If `buf` is too small for the given dimensions this function will panic
211#[deny(clippy::arithmetic_side_effects)]
212pub fn infer_i41x<S: AnySlice>(
213    buf: S,
214    width: usize,
215    height: usize,
216    strides: Option<&[usize]>,
217) -> [S; 3] {
218    infer_impl(I41X_PLANES, buf, width, height, strides)
219}
220
221/// Infer the planes for a full NV12 image using the given dimensions
222///
223/// # Panics
224///
225/// If `buf` is too small for the given dimensions this function will panic
226#[deny(clippy::arithmetic_side_effects)]
227pub fn infer_nv12<S: AnySlice>(
228    buf: S,
229    width: usize,
230    height: usize,
231    strides: Option<&[usize]>,
232) -> [S; 2] {
233    infer_impl(NV12_PLANES, buf, width, height, strides)
234}
235
236/// Infer the planes for a full P010 / P012 image using the given dimensions
237///
238/// # Panics
239///
240/// If `buf` is too small for the given dimensions this function will panic
241#[deny(clippy::arithmetic_side_effects)]
242pub fn infer_p01x<S: AnySlice>(
243    buf: S,
244    width: usize,
245    height: usize,
246    strides: Option<&[usize]>,
247) -> [S; 2] {
248    infer_impl(P01X_PLANES, buf, width, height, strides)
249}
250
251/// Helper trait implemented on &[T] and &mut [T]
252#[diagnostic::on_unimplemented(message = "AnySlice is only implemented for &[T] and &mut [T].\n\
253               When using or Vec<T> or similar try .as_slice() or .as_mut_slice()")]
254pub trait AnySlice: sealed::Sealed + Default + Sized {
255    fn slice_len(&self) -> usize;
256    fn slice_split_at(self, at: usize) -> (Self, Self);
257}
258
259mod sealed {
260    pub trait Sealed {}
261    impl<T> Sealed for &[T] {}
262    impl<T> Sealed for &mut [T] {}
263    impl<T> Sealed for Vec<T> {}
264}
265
266impl<T> AnySlice for &[T] {
267    fn slice_len(&self) -> usize {
268        self.len()
269    }
270
271    fn slice_split_at(self, at: usize) -> (Self, Self) {
272        self.split_at(at)
273    }
274}
275
276impl<T> AnySlice for &mut [T] {
277    fn slice_len(&self) -> usize {
278        self.len()
279    }
280
281    fn slice_split_at(self, at: usize) -> (Self, Self) {
282        self.split_at_mut(at)
283    }
284}