mux: reassemble MPEG-2 access units via the shared AuAssembler

The MPEG-2 parser hand-rolled its own PES reassembly — a byte buffer plus
parallel PTS / source / discontinuity mark queues keyed by absolute
offset — duplicating what AuAssembler already does for H.264/HEVC/VC-1.

Add a Mode::Mpeg2 to AuAssembler (picture 0x00 with preceding sequence
0xB3 / GOP 0xB8 headers — the same headers-precede-picture shape as the
VC-1 mode) and have the MPEG-2 parser own one via AuAssembler::mpeg2().
parse() now feeds fragments to the assembler and processes each complete
access unit; the buffer, base offset, and three mark queues are gone. The
GOP-buffered temporal_reference reorder and PTS origin-locking are
unchanged. The parser's external contract is unchanged, so all existing
MPEG-2 parser tests pass as-is; new AuAssembler tests cover the MPEG-2
boundary rule directly.

src/mux/au_assembly.rs

@@ -7,14 +7,15 @@
 //! chops the elementary stream into fixed-size PES fragments with no AU
 //! alignment, and only the first fragment of an AU carries a PTS. So a parser
 //! that assumes one-AU-per-PES (h264/hevc/vc1, written against TS) mis-frames a
-//! program stream, while `mpeg2` — the DVD/PS codec — has always reassembled
-//! across PES in its own parser.
+//! program stream, while `mpeg2` — the DVD/PS codec — must reassemble across PES.
 //!
-//! [`AuAssembler`] is that reassembly, factored out so the h264/hevc/vc1 parsers
-//! can do what `mpeg2` already does without hand-rolling the buffer three times.
-//! It buffers PES-fragment bytes and emits one AU per codec AU boundary, carrying
-//! the AU-start timing/source forward. Since the boundary is a codec start code,
-//! it lives with the codec parser (which picks the marker); only the generic
+//! [`AuAssembler`] is that reassembly, factored out so EVERY program-stream video
+//! parser shares one implementation instead of hand-rolling the buffer. The
+//! h264/hevc/vc1 parsers ([`Mode::StartCode`] / [`Mode::Vc1`]) and the MPEG-2
+//! parser ([`Mode::Mpeg2`], via [`AuAssembler::mpeg2`]) all drive it. It buffers
+//! PES-fragment bytes and emits one AU per codec AU boundary, carrying the
+//! AU-start timing/source forward. Since the boundary is a codec start code, it
+//! lives with the codec parser (which picks the marker); only the generic
 //! buffering + timing-carry is shared here.
 //!
 //! This is *inside* the parser, not a pipeline stage: the pipeline stays
@@ -46,6 +47,11 @@ const VC1_FRAME: u8 = 0x0D; // coded picture
 const VC1_ENTRY: u8 = 0x0E; // entry-point header
 const VC1_SEQ: u8 = 0x0F; // sequence header
 
+/// MPEG-2 (ISO/IEC 13818-2) start-code suffixes, `00 00 01 <type>`.
+const MP2_PICTURE: u8 = 0x00; // picture_start_code
+const MP2_SEQ: u8 = 0xB3; // sequence_header_code
+const MP2_GOP: u8 = 0xB8; // group_start_code
+
 /// How a stream's fragments become AU-complete units.
 #[derive(Clone, Copy)]
 enum Mode {
@@ -60,6 +66,13 @@ enum Mode {
     /// I-frame of its headers. The boundary is instead the next `0x0F`/`0x0E`/`0x0D`
     /// start code that follows a frame already seen in the current AU.
     Vc1,
+    /// MPEG-2 access unit: `[sequence header?][GOP header?][picture][slices…]`.
+    /// Structurally identical to [`Mode::Vc1`] — the sequence (`0xB3`) and GOP
+    /// (`0xB8`) headers precede the picture (`0x00`) they introduce, so the
+    /// boundary is the next picture / sequence / GOP start code that follows a
+    /// picture already seen. Slice (`0x01..=0xAF`), extension (`0xB5`),
+    /// user-data (`0xB2`) and sequence-end (`0xB7`) codes are NOT boundaries.
+    Mpeg2,
     /// The codec self-frames (MPEG-2 reassembles in its own parser; audio resyncs
     /// on syncwords), so each fragment passes straight through as one unit. Lets
     /// the caller run EVERY stream through an assembler with no per-codec branch.
@@ -112,6 +125,20 @@ impl AuAssembler {
         }
     }
 
+    /// An assembler that reassembles MPEG-2 access units. The MPEG-2 parser owns
+    /// one of these directly (rather than hand-rolling the buffer): the demux
+    /// layer runs MPEG-2 through [`Mode::Passthrough`] and hands each fragment to
+    /// the parser, which feeds them here to be reframed on picture boundaries.
+    pub(crate) fn mpeg2() -> Self {
+        Self {
+            mode: Mode::Mpeg2,
+            buf: Vec::with_capacity(128 * 1024),
+            base: 0,
+            marks: VecDeque::new(),
+            disc_marks: VecDeque::new(),
+        }
+    }
+
     /// Feed one PES fragment; return every AU that is now complete.
     pub(crate) fn push(
         &mut self,
@@ -251,6 +278,8 @@ fn au_opener(mode: Mode, buf: &[u8]) -> Option<usize> {
         Mode::StartCode(marker) => find_start_code(buf, 0, marker),
         // Any of the three AU-opening BDU types opens a VC-1 access unit.
         Mode::Vc1 => find_vc1_start(buf, 0),
+        // A sequence header, GOP header, or picture opens an MPEG-2 access unit.
+        Mode::Mpeg2 => find_mpeg2_start(buf, 0),
         Mode::Passthrough => None,
     }
 }
@@ -262,6 +291,7 @@ fn au_boundary(mode: Mode, buf: &[u8]) -> Option<usize> {
         // AU ends at the next delimiter; skip the opening one at buf[0].
         Mode::StartCode(marker) => find_start_code(buf, 4, marker),
         Mode::Vc1 => find_vc1_au_end(buf),
+        Mode::Mpeg2 => find_mpeg2_au_end(buf),
         Mode::Passthrough => None,
     }
 }
@@ -327,6 +357,56 @@ fn find_vc1_au_end(buf: &[u8]) -> Option<usize> {
     None
 }
 
+/// Find the next MPEG-2 AU-opening start code (`00 00 01` followed by a picture,
+/// sequence header, or GOP header) at or after `from`.
+fn find_mpeg2_start(buf: &[u8], from: usize) -> Option<usize> {
+    let mut i = from;
+    while i + 4 <= buf.len() {
+        if buf[i] == 0
+            && buf[i + 1] == 0
+            && buf[i + 2] == 1
+            && matches!(buf[i + 3], MP2_PICTURE | MP2_SEQ | MP2_GOP)
+        {
+            return Some(i);
+        }
+        i += 1;
+    }
+    None
+}
+
+/// End offset of the MPEG-2 access unit that opens at `buf[0]`: the next picture
+/// / sequence / GOP start code that appears *after* this AU already contains a
+/// picture (`0x00`). Returns `None` while the AU is still open (no picture yet,
+/// or no following boundary buffered). A leading sequence/GOP header thus stays
+/// attached to the picture it introduces. Slice / extension / user-data /
+/// sequence-end codes are skipped — they belong to the current AU.
+fn find_mpeg2_au_end(buf: &[u8]) -> Option<usize> {
+    let mut seen_picture = false;
+    let mut i = 0usize;
+    while i + 4 <= buf.len() {
+        if buf[i] == 0 && buf[i + 1] == 0 && buf[i + 2] == 1 {
+            match buf[i + 3] {
+                MP2_PICTURE => {
+                    if i > 0 && seen_picture {
+                        return Some(i);
+                    }
+                    seen_picture = true;
+                }
+                MP2_SEQ | MP2_GOP => {
+                    if i > 0 && seen_picture {
+                        return Some(i);
+                    }
+                }
+                _ => {}
+            }
+            i += 4;
+        } else {
+            i += 1;
+        }
+    }
+    None
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -501,6 +581,62 @@ mod tests {
         assert_eq!(done[0].pts, Some(500));
     }
 
+    // ── MPEG-2 AU grouping ────────────────────────────────────────────────
+
+    #[test]
+    fn mpeg2_keeps_seq_and_gop_headers_with_their_picture() {
+        // A GOP-opening AU is [seq 0xB3][gop 0xB8][picture 0x00][slices]; the next
+        // picture (no headers) is its own AU. The seq/GOP headers must stay with
+        // the picture they introduce, not glue onto the previous AU.
+        let mut a = AuAssembler::mpeg2();
+        let mut gop = bdu(MP2_SEQ, 0xAA, 10);
+        gop.extend(bdu(MP2_GOP, 0xBB, 8));
+        gop.extend(bdu(MP2_PICTURE, 0xCC, 20)); // picture + slice bytes
+        let pic2 = bdu(MP2_PICTURE, 0xDD, 15);
+
+        assert!(a.push(&gop, Some(9000), None, None, false).is_empty());
+        let out = a.push(&pic2, Some(9376), None, None, false);
+        assert_eq!(
+            out.len(),
+            1,
+            "first AU completes at the next picture boundary"
+        );
+        assert_eq!(out[0].data, gop, "AU retains seq + GOP + picture");
+        assert_eq!(out[0].pts, Some(9000));
+
+        let tail = a.flush();
+        assert_eq!(tail.len(), 1);
+        assert_eq!(tail[0].data, pic2, "second picture is its own AU");
+        assert_eq!(tail[0].pts, Some(9376));
+    }
+
+    #[test]
+    fn mpeg2_slice_codes_are_not_au_boundaries() {
+        // Slice start codes (0x01..=0xAF) inside a picture must not split the AU.
+        let mut a = AuAssembler::mpeg2();
+        let mut pic = bdu(MP2_PICTURE, 0x11, 4);
+        pic.extend(bdu(0x01, 0x22, 10)); // slice 1
+        pic.extend(bdu(0xAF, 0x33, 10)); // slice 175 (max slice code)
+        let next = bdu(MP2_PICTURE, 0x44, 4); // opening boundary of the next AU
+        let out = a.push(&[pic.clone(), next].concat(), Some(1), None, None, false);
+        assert_eq!(out.len(), 1, "slices stay inside the one picture AU");
+        assert_eq!(out[0].data, pic, "AU spans the picture and all its slices");
+    }
+
+    #[test]
+    fn mpeg2_reassembles_one_picture_split_across_fragments() {
+        // A picture split across three PES fragments; only the first carries a PTS.
+        let mut a = AuAssembler::mpeg2();
+        let full = bdu(MP2_PICTURE, 0xEE, 100);
+        assert!(a.push(&full[..40], Some(500), None, None, false).is_empty());
+        assert!(a.push(&full[40..80], None, None, None, false).is_empty());
+        assert!(a.push(&full[80..], None, None, None, false).is_empty());
+        let out = a.flush();
+        assert_eq!(out.len(), 1);
+        assert_eq!(out[0].pts, Some(500), "AU carries its START pts");
+        assert_eq!(out[0].data, full);
+    }
+
     #[test]
     fn over_cap_without_boundary_force_flushes() {
         let mut a = AuAssembler::for_codec(Codec::H264);