yuv420: Fix plane packing

There is no way to guess width and height from the length of the buffer,
except for very specific cases (e.g. a 6x4 image, where the length of the
U and V planes matches the width by accidence).

There was another bug where the i (row) and j (column) iterators were
zipped, which lead to them increasing in lockstep. Instead, we want a
nested loop.

Signed-off-by: Christopher N. Hesse <[email protected]>

ffimage-yuv/benches/convert.rs

@@ -56,14 +56,14 @@ pub fn yuv420p_to_rgb(c: &mut Criterion) {
     let resolutions = [(640, 480), (1280, 720)];
 
     for res in resolutions {
-        let yuv420p = vec![10; ((res.0 * res.1) as f32 * 1.5) as usize];
+        let yuv420p = vec![10; res.0 * res.1 * 3 / 2];
         let mut rgb = vec![10; res.0 * res.1 * 3];
 
         c.bench_function(
             &format!("Yuv420p[u8] -> Rgb[u8] ({}x{})", res.0, res.1),
             |b| {
                 b.iter(|| {
-                    Yuv420p::pack(&yuv420p)
+                    Yuv420p::pack(&yuv420p, 640, 480)
                         .into_iter()
                         .colorconvert::<Rgb<u8>>()
                         .bytes()

ffimage-yuv/src/yuv420.rs

@@ -9,17 +9,25 @@ pub struct Yuv420p;
 
 impl Yuv420p {
     /// Returns packed Yuv444 color samples from a given slice.
-    pub fn pack<'a, T>(buf: &'a [T]) -> impl IntoIterator<Item = Yuv<T>> + 'a
+    pub fn pack<'a, T>(
+        buf: &'a [T],
+        width: u32,
+        height: u32,
+    ) -> impl IntoIterator<Item = Yuv<T>> + 'a
     where
         T: Copy,
     {
-        Self::pack_bytes(buf)
+        Self::pack_bytes(buf, width, height)
             .into_iter()
             .map(|chunk| Yuv::<T>::from(chunk))
     }
 
     /// Returns packed Yuv444 color samples as byte chunks from a given slice.
-    pub fn pack_bytes<'a, T>(buf: &'a [T]) -> impl IntoIterator<Item = [T; 3]> + 'a
+    pub fn pack_bytes<'a, T>(
+        buf: &'a [T],
+        width: u32,
+        height: u32,
+    ) -> impl IntoIterator<Item = [T; 3]> + 'a
     where
         T: Copy,
     {
@@ -28,41 +36,43 @@ impl Yuv420p {
         // buf must be divisible in 3 parts: 2/3 Luma (y) samples, 1/3 Chroma (u + v) samples
         assert_eq!(buf.len() % 3, 0);
 
-        let pixels = (buf.len() / 3) * 2;
-        let width = pixels / 4;
+        let luma_count = buf.len() / 3 * 2;
+        let chroma_count = buf.len() / 3 / 2;
+        assert_eq!(luma_count + chroma_count * 2, buf.len());
 
-        let y = &buf[0..pixels];
-        let u = &buf[pixels..(pixels + width)];
-        let v = &buf[(pixels + width)..(pixels + width * 2)];
+        let y = &buf[0..luma_count];
+        let u = &buf[luma_count..(luma_count + chroma_count)];
+        let v = &buf[(luma_count + chroma_count)..(luma_count + chroma_count * 2)];
+        assert_eq!(y.len(), (u.len() + v.len()) * 2);
+        assert_eq!(y.len() + u.len() + v.len(), buf.len());
 
-        Yuv420p::pack_planes(y, u, v)
+        Yuv420p::pack_planes(y, u, v, width, height)
     }
 
     /// Returns packed Yuv444 color samples as byte chunks from YUV planes.
     pub fn pack_planes<'a, T>(
         y: &'a [T],
         u: &'a [T],
         v: &'a [T],
+        width: u32,
+        height: u32,
     ) -> impl IntoIterator<Item = [T; 3]> + 'a
     where
         T: Copy,
     {
         // YUV420 has 2x2 blocks of Luma (y) samples
         assert_eq!(y.len() % 4, 0);
-        let width = y.len() / 4;
-        let height = y.len() / width;
-        assert_eq!(y.len(), width * height);
-        assert_eq!(u.len(), width);
-        assert_eq!(v.len(), width);
 
-        (0..height)
+        (0..height as usize)
             .into_iter()
-            .zip((0..width).into_iter())
-            .map(move |(i, j)| {
-                let y_idx = i * width + j;
-                let uv_idx = i / 2 * width / 2 + j / 2;
-                [y[y_idx], u[uv_idx], v[uv_idx]]
+            .map(move |i| {
+                (0..width as usize).into_iter().map(move |j| {
+                    let y_idx = i * width as usize + j;
+                    let uv_idx = i / 2 * width as usize / 2 + j / 2;
+                    [y[y_idx], u[uv_idx], v[uv_idx]]
+                })
             })
+            .flatten()
     }
 
     /// Unpacks packed Yuv444 color samples into Y, U, V planes.