mxf: faster decode_pcm, AES3 read_pcm no longer decodes the whole element per call

crates/mxf/src/essence.rs

@@ -201,6 +201,38 @@ pub fn aes3_element_samples(v: &[u8]) -> usize {
     declared.min((v.len() - 4) / 32)
 }
 
+/// De-interleave `bps`-byte little-endian samples (`B` = `bps`), one pass per channel.
+fn decode_planes<const B: usize>(v: &[u8], frame_bytes: usize, out: &mut [Vec<f32>], scale: f32) {
+    for (c, o) in out.iter_mut().enumerate() {
+        let off = c * B;
+        o.extend(v.chunks_exact(frame_bytes).map(|f| {
+            let mut s = [0u8; B];
+            s.copy_from_slice(&f[off..off + B]);
+            let x = match B {
+                1 => (s[0] as i32 - 128) << 24,
+                2 => i32::from_le_bytes([0, 0, s[0], s[1]]),
+                3 => i32::from_le_bytes([0, s[0], s[1], s[2]]),
+                _ => i32::from_le_bytes([s[0], s[1], s[2], s[3]]),
+            };
+            (x >> (32 - 8 * B)) as f32 * scale
+        }));
+    }
+}
+
+/// Decode ST 331 sample frames (without the element header; 32 bytes each, the trailing partial
+/// frame is ignored) into `channels` planar channels, at most eight.
+pub(crate) fn decode_aes3_frames(frames: &[u8], channels: usize) -> Vec<Vec<f32>> {
+    let mut out: Vec<Vec<f32>> = (0..channels).map(|_| Vec::with_capacity(frames.len() / 32)).collect();
+    for f in frames.as_chunks::<32>().0 {
+        for (c, o) in out.iter_mut().enumerate().take(8) {
+            let w = u32::from_le_bytes([f[c * 4], f[c * 4 + 1], f[c * 4 + 2], f[c * 4 + 3]]);
+            let s = (((w >> 4) & 0xFF_FFFF) << 8) as i32 >> 8;
+            o.push(s as f32 / 8_388_608.0);
+        }
+    }
+    out
+}
+
 /// Decode PCM chunk bytes into planar f32 (`channels` channels).
 ///
 /// `Pcm`: interleaved little-endian signed integers of `bits` (8-bit is unsigned, per WAVE),
@@ -219,20 +251,14 @@ pub fn decode_pcm(v: &[u8], format: SoundFormat, channels: usize, bits: u32, fra
             }
             // `frame_bytes >= bps * ch`, so the output is bounded by the bytes actually supplied.
             let n = v.len() / frame_bytes;
-            let mut out = vec![Vec::with_capacity(n); ch];
             let scale = 1.0 / (1u64 << (8 * bps - 1)) as f32;
-            for i in 0..n {
-                let f = &v[i * frame_bytes..];
-                for (c, o) in out.iter_mut().enumerate() {
-                    let s = &f[c * bps..c * bps + bps];
-                    let x = match bps {
-                        1 => (s[0] as i32 - 128) << 24,
-                        2 => i32::from_le_bytes([0, 0, s[0], s[1]]),
-                        3 => i32::from_le_bytes([0, s[0], s[1], s[2]]),
-                        _ => i32::from_le_bytes([s[0], s[1], s[2], s[3]]),
-                    };
-                    o.push((x >> (32 - 8 * bps)) as f32 * scale);
-                }
+            // `vec![Vec::with_capacity(n); ch]` would clone empty vectors without the capacity.
+            let mut out: Vec<Vec<f32>> = (0..ch).map(|_| Vec::with_capacity(n)).collect();
+            match bps {
+                1 => decode_planes::<1>(v, frame_bytes, &mut out, scale),
+                2 => decode_planes::<2>(v, frame_bytes, &mut out, scale),
+                3 => decode_planes::<3>(v, frame_bytes, &mut out, scale),
+                _ => decode_planes::<4>(v, frame_bytes, &mut out, scale),
             }
             Ok(out)
         }
@@ -243,16 +269,7 @@ pub fn decode_pcm(v: &[u8], format: SoundFormat, channels: usize, bits: u32, fra
                 return Err(crate::Error::Invalid("an AES3 element has at most eight channels".into()));
             }
             let n = aes3_element_samples(v);
-            let mut out = vec![Vec::with_capacity(n); ch];
-            for i in 0..n {
-                for (c, o) in out.iter_mut().enumerate().take(8) {
-                    let p = 4 + (i * 8 + c) * 4;
-                    let w = u32::from_le_bytes([v[p], v[p + 1], v[p + 2], v[p + 3]]);
-                    let s = (((w >> 4) & 0xFF_FFFF) << 8) as i32 >> 8;
-                    o.push(s as f32 / 8_388_608.0);
-                }
-            }
-            Ok(out)
+            Ok(decode_aes3_frames(&v[4.min(v.len())..][..n * 32], ch))
         }
         SoundFormat::Other => Err(crate::Error::Unsupported("sound coding".into())),
     }
@@ -324,4 +341,48 @@ mod tests {
         // truncated: no panic, no samples
         assert!(decode_pcm(&e[..10], SoundFormat::Aes3Element, 2, 24, 0).unwrap()[0].is_empty());
     }
+
+    /// The per-sample loop `decode_pcm` used before the per-width passes.
+    fn decode_pcm_reference(v: &[u8], ch: usize, bits: u32, frame_bytes: usize) -> Vec<Vec<f32>> {
+        let bps = (bits as usize).div_ceil(8);
+        let n = v.len() / frame_bytes;
+        let mut out = vec![Vec::with_capacity(n); ch];
+        let scale = 1.0 / (1u64 << (8 * bps - 1)) as f32;
+        for i in 0..n {
+            let f = &v[i * frame_bytes..];
+            for (c, o) in out.iter_mut().enumerate() {
+                let s = &f[c * bps..c * bps + bps];
+                let x = match bps {
+                    1 => (s[0] as i32 - 128) << 24,
+                    2 => i32::from_le_bytes([0, 0, s[0], s[1]]),
+                    3 => i32::from_le_bytes([0, s[0], s[1], s[2]]),
+                    _ => i32::from_le_bytes([s[0], s[1], s[2], s[3]]),
+                };
+                o.push((x >> (32 - 8 * bps)) as f32 * scale);
+            }
+        }
+        out
+    }
+
+    #[test]
+    fn decode_pcm_matches_the_per_sample_reference() {
+        let mut seed = 0x1234_5678_u32;
+        let v: Vec<u8> = (0..4099)
+            .map(|_| {
+                seed = seed.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
+                (seed >> 24) as u8
+            })
+            .collect();
+        for (bits, ch, pad) in [(8, 1, 0), (8, 2, 2), (16, 1, 0), (16, 2, 0), (16, 6, 3), (24, 1, 0), (24, 2, 0), (24, 3, 5), (32, 2, 0), (32, 1, 4)] {
+            let bps = (bits as usize).div_ceil(8);
+            let fb = bps * ch + pad;
+            let got = decode_pcm(&v, SoundFormat::Pcm, ch, bits, fb).unwrap();
+            let want = decode_pcm_reference(&v, ch, bits, fb);
+            assert_eq!(got.len(), want.len());
+            for (g, w) in got.iter().zip(&want) {
+                assert!(g.iter().map(|x| x.to_bits()).eq(w.iter().map(|x| x.to_bits())), "{bits}-bit, {ch} channels, {pad} pad");
+                assert_eq!(g.capacity(), v.len() / fb);
+            }
+        }
+    }
 }

crates/mxf/src/file.rs

@@ -3,7 +3,7 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use crate::essence::{Codec, PictureInfo, SoundFormat, SoundInfo, decode_pcm, identify_picture, sniff_picture};
+use crate::essence::{Codec, PictureInfo, SoundFormat, SoundInfo, decode_aes3_frames, decode_pcm, identify_picture, sniff_picture};
 use crate::index::{IndexEntry, IndexSegment};
 use crate::klv::{KeyClass, Ul, ber_length, classify, find_header_partition};
 use crate::meta::{Metadata, Primer, SUB_DESCRIPTORS_ITEM, Set, set_type as st, tag};
@@ -1063,13 +1063,28 @@ impl MxfFile {
                         decode_pcm(&v, SoundFormat::Pcm, ch, si.bits, fb)?
                     }
                     SoundFormat::Aes3Element => {
-                        let size = c.size.min(src.len().saturating_sub(c.offset));
-                        let size = usize::try_from(size).map_err(|_| Error::Invalid("AES3 chunk exceeds the address space".into()))?;
-                        let mut v = vec![0u8; size];
-                        src.read_at(c.offset, &mut v)?;
-                        let all = decode_pcm(&v, SoundFormat::Aes3Element, ch, si.bits, 0)?;
-                        let s0 = (a - c.first_sample) as usize;
-                        all.into_iter().map(|p| p.get(s0..).map(<[f32]>::to_vec).unwrap_or_default()).collect()
+                        // Read only the requested sample frames (32 bytes each behind the 4-byte
+                        // element header), not the whole element: a clip-wrapped or hostile
+                        // element would otherwise be read and decoded in full on every call.
+                        if ch > 8 {
+                            return Err(Error::Invalid("an AES3 element has at most eight channels".into()));
+                        }
+                        let avail = c.size.min(src.len().saturating_sub(c.offset));
+                        let mut head = [0u8; 4];
+                        let total = if avail >= 4 {
+                            src.read_at(c.offset, &mut head)?;
+                            // same bound as `aes3_element_samples`
+                            u64::from(u16::from_le_bytes([head[1], head[2]])).min((avail - 4) / 32)
+                        } else {
+                            0
+                        };
+                        let s0 = a - c.first_sample;
+                        let n = (b - a).min(total.saturating_sub(s0));
+                        let mut v = vec![0u8; n as usize * 32];
+                        if n > 0 {
+                            src.read_at(c.offset + 4 + s0 * 32, &mut v)?;
+                        }
+                        decode_aes3_frames(&v, ch)
                     }
                     SoundFormat::Other => return Err(Error::Unsupported("sound essence coding".into())),
                 };
@@ -1109,3 +1124,101 @@ mod index_bounds_tests {
         assert_eq!(entries.keys().copied().collect::<Vec<_>>(), vec![0, 1]);
     }
 }
+
+#[cfg(test)]
+mod aes3_read_tests {
+    use std::cell::Cell;
+
+    use super::*;
+    use crate::source::ByteSource;
+
+    /// A source that remembers its largest single read.
+    struct Counting<'a> {
+        data: &'a [u8],
+        biggest: Cell<usize>,
+    }
+
+    impl ByteSource for Counting<'_> {
+        fn len(&self) -> u64 {
+            self.data.len() as u64
+        }
+        fn read_at(&self, offset: u64, buf: &mut [u8]) -> std::io::Result<()> {
+            self.biggest.set(self.biggest.get().max(buf.len()));
+            self.data.read_at(offset, buf)
+        }
+    }
+
+    /// One clip-wrapped AES3 element of `declared` sample frames in a file that holds `stored`.
+    fn aes3_file(declared: u16, stored: usize, channels: u32) -> (MxfFile, Vec<u8>) {
+        let mut data = vec![0x00, declared as u8, (declared >> 8) as u8, 0xFF];
+        for i in 0..stored * 8 {
+            let s = ((i as i32 * 7919) % 1_000_000) - 500_000;
+            data.extend_from_slice(&((((s as u32) & 0xFF_FFFF) << 4) | (i as u32 & 15)).to_le_bytes());
+        }
+        let samples = (data.len() as u64 - 4) / 32;
+        let track = EssenceTrack {
+            track_id: 1,
+            track_number: 0,
+            kind: TrackKind::Sound,
+            name: None,
+            edit_rate: Rational::default(),
+            origin: 0,
+            start_position: 0,
+            duration: None,
+            codec: Codec::Aes3Element,
+            wrapping: Wrapping::Clip,
+            essence_container: None,
+            picture: None,
+            sound: Some(SoundInfo { channels, bits: 24, ..Default::default() }),
+            sound_format: SoundFormat::Aes3Element,
+            body_sid: 1,
+            index_sid: 0,
+            samples: Vec::new(),
+            display_order: Vec::new(),
+            chunks: vec![Chunk { offset: 0, size: data.len() as u64, first_sample: 0, samples }],
+            temporal_offsets: false,
+            needs_reorder: false,
+            indexed: false,
+        };
+        let file = MxfFile {
+            run_in: 0,
+            operational_pattern: OperationalPattern::Atom,
+            operational_pattern_label: Ul::default(),
+            partitions: Vec::new(),
+            index_segments: Vec::new(),
+            tracks: vec![track],
+            timecode: None,
+            material_package_name: None,
+            product_name: None,
+            company_name: None,
+            metadata_partition: 0,
+            warnings: Vec::new(),
+            file_size: data.len() as u64,
+        };
+        (file, data)
+    }
+
+    #[test]
+    fn aes3_reads_decode_only_the_requested_frames() {
+        for (declared, stored) in [(5000, 5000), (3000, 5000)] {
+            let (file, data) = aes3_file(declared, stored, 2);
+            let whole = decode_pcm(&data, SoundFormat::Aes3Element, 2, 24, 0).unwrap();
+            assert_eq!(whole[0].len(), declared as usize);
+            for (start, frames) in [(0u64, 1920usize), (1920, 1920), (2999, 4), (2990, 2000), (4990, 100), (6000, 10)] {
+                let src = Counting { data: &data, biggest: Cell::new(0) };
+                let got = file.read_pcm(&src, 0, start, frames).unwrap();
+                for (c, g) in got.iter().enumerate() {
+                    let want: Vec<f32> = (0..frames).map(|k| whole[c].get(start as usize + k).copied().unwrap_or(0.0)).collect();
+                    assert_eq!(g, &want, "declared {declared}, start {start}, frames {frames}, channel {c}");
+                }
+                assert!(src.biggest.get() <= frames * 32, "read {} bytes for {frames} frames", src.biggest.get());
+            }
+        }
+    }
+
+    #[test]
+    fn aes3_read_rejects_more_than_eight_channels() {
+        let (file, data) = aes3_file(10, 10, 9);
+        assert!(file.read_pcm(&data, 0, 0, 4).is_err());
+    }
+}