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
+//! program stream, while `mpeg2` — the DVD/PS codec — must reassemble across PES.
-//! across PES in its own parser.
 //!
-//! [`AuAssembler`] is that reassembly, factored out so the h264/hevc/vc1 parsers
+//! [`AuAssembler`] is that reassembly, factored out so EVERY program-stream video
-//! can do what `mpeg2` already does without hand-rolling the buffer three times.
+//! parser shares one implementation instead of hand-rolling the buffer. The
-//! It buffers PES-fragment bytes and emits one AU per codec AU boundary, carrying
+//! h264/hevc/vc1 parsers ([`Mode::StartCode`] / [`Mode::Vc1`]) and the MPEG-2
-//! the AU-start timing/source forward. Since the boundary is a codec start code,
+//! parser ([`Mode::Mpeg2`], via [`AuAssembler::mpeg2`]) all drive it. It buffers
-//! it lives with the codec parser (which picks the marker); only the generic
+//! 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);

src/mux/codec/mpeg2.rs

@@ -26,13 +26,10 @@
 //! - Extension (seq/pic):00 00 01 B5
 //! - GOP header:         00 00 01 B8
 
-use std::collections::VecDeque;
-
 use super::coding::{CodingType, Mpeg2Coding, PictureInfo};
 use super::startcode::find_start_code;
 use super::{CodecParser, Frame, pts_to_ns};
 use crate::mux::ts::PesPacket;
-use crate::pes::SourcePos;
 
 /// Sequence header start code suffix.
 const SEQ_HEADER_CODE: u8 = 0xB3;
@@ -94,20 +91,11 @@ pub struct Mpeg2Parser {
     /// Raw bytes of the last seen sequence header (+ sequence extension if
     /// present), captured for MKV codecPrivate.
     seq_header: Option<Vec<u8>>,
-    /// Unemitted elementary-stream bytes: the in-progress access unit plus any
+    /// Reassembles PES fragments into complete access units (one coded picture
-    /// lookahead needed to detect the next AU boundary.
+    /// with its leading sequence/GOP headers) and carries each AU's start
-    buf: Vec<u8>,
+    /// timing / source / discontinuity forward — the shared machinery the
-    /// Absolute ES byte offset of `buf[0]`. Used to associate PES PTS marks
+    /// H.264/HEVC/VC-1 parsers also use, in its MPEG-2 mode.
-    /// (recorded by absolute offset) with the access units they belong to.
+    au_asm: crate::mux::au_assembly::AuAssembler,
-    base_offset: u64,
-    /// `(absolute ES offset of a PES's first byte, PTS in ns)` for every PES
-    /// that carried a timestamp, in ascending offset order.
-    pts_marks: VecDeque<(u64, i64)>,
-    /// `(absolute ES offset of a PES's first byte, SourcePos)` for every PES
-    /// that carried byte-exact provenance, parallel to `pts_marks` and drained
-    /// by the SAME mark-drain invariant. Attaches the source position to each
-    /// access unit so the index carries it — never reconstructed.
-    source_marks: VecDeque<(u64, SourcePos)>,
     /// Full-frame presentation interval (ns) at the sequence-header display rate
     /// (`1/frame_rate`). The field period is half this. Per-frame durations are
     /// `nb_fields × field_period`, so 2:3-telecined frames alternate 2- and
@@ -129,15 +117,6 @@ pub struct Mpeg2Parser {
     /// each GOP's first PES PTS so video stays in sync with the PES-timestamped
     /// audio. None until the first PES timestamp is seen.
     origin_pts_ns: Option<i64>,
-    /// B1: absolute ES offsets at which a concealed/lost-gap PES began, parallel
-    /// to `pts_marks`/`source_marks` and drained by the SAME mark-drain invariant.
-    /// MPEG-2 emits whole GOPs asynchronously, so a per-PES flag can't ride
-    /// through to the right frame (the PES that carries the gap completes the
-    /// PREVIOUS picture); associating by OFFSET instead stamps `discontinuity` on
-    /// the access unit whose own bytes begin after the gap — the first post-gap
-    /// picture — surviving GOP buffering + temporal reorder. The consumer's
-    /// ResyncGate then arms at that exact picture, mid-GOP if need be.
-    disc_marks: VecDeque<u64>,
 }
 
 /// One coded picture buffered awaiting its GOP's completion (see `gop_buf`).
@@ -165,16 +144,12 @@ impl Mpeg2Parser {
     pub fn new() -> Self {
         Self {
             seq_header: None,
-            buf: Vec::with_capacity(128 * 1024),
+            au_asm: crate::mux::au_assembly::AuAssembler::mpeg2(),
-            base_offset: 0,
-            pts_marks: VecDeque::new(),
-            source_marks: VecDeque::new(),
             frame_duration_ns: 0,
             progressive_sequence: false,
             gop_buf: Vec::new(),
             emitted_fields: 0,
             origin_pts_ns: None,
-            disc_marks: VecDeque::new(),
         }
     }
 
@@ -202,192 +177,95 @@ impl Mpeg2Parser {
     /// Drain every complete access unit from `buf`, returning one Frame each.
     /// When `force` is true (EOF flush, or buffer-cap backstop) the trailing
     /// in-progress access unit is emitted even without a following boundary.
-    fn drain_complete_aus(&mut self, force: bool) -> Vec<Frame> {
+    /// Process one reassembled access unit (from [`AuAssembler`]): decode its
-        let mut out = Vec::new();
+    /// per-picture coding info, capture a new sequence header, and buffer the
-        loop {
+    /// picture into the current GOP for display-order timestamping. The AU's
-            // An access unit must contain a coded picture; without one there is
+    /// timing / source / discontinuity were already attributed by the assembler.
-            // nothing to emit yet (leading sequence/GOP headers wait for it).
+    fn process_au(&mut self, au: crate::mux::au_assembly::AssembledAu, out: &mut Vec<Frame>) {
-            let Some(pic) = find_code(&self.buf, 0, PICTURE_CODE) else {
+        let data = au.data;
-                // No coded picture in an over-cap buffer means we are
+        // An access unit must contain a coded picture; a fragment that assembled
-                // accumulating unparseable data (a stream with no picture
+        // without one (only headers, or truncated at EOF) yields nothing.
-                // start codes). Drop all but a 3-byte tail — enough to catch a
+        let Some(pic) = find_code(&data, 0, PICTURE_CODE) else {
-                // start-code prefix straddling the boundary — and advance the
+            return;
-                // absolute offset so the PES-mark invariant holds. Mirrors the
+        };
-                // post-picture buffer backstop in the AU-boundary search below.
+        let end = data.len();
-                if self.buf.len() > MAX_AU_BUFFER {
+        // Capture a sequence header for codecPrivate; a new one replaces the
-                    let drop = self.buf.len() - 3;
+        // stored value and re-locks the frame duration.
-                    self.base_offset += drop as u64;
+        if let Some(h) = extract_seq_header(&data) {
-                    self.buf.drain(..drop);
+            self.progressive_sequence = parse_progressive_sequence(&h);
-                    let cutoff = self.base_offset;
+            self.seq_header = Some(h);
-                    while let Some(&(off, _)) = self.pts_marks.front() {
+            if let Some((num, den)) = self.frame_rate() {
-                        if off < cutoff {
+                if num > 0 {
-                            self.pts_marks.pop_front();
+                    self.frame_duration_ns = 1_000_000_000i64 * den as i64 / num as i64;
-                        } else {
-                            break;
-                        }
-                    }
-                    while let Some(&(off, _)) = self.source_marks.front() {
-                        if off < cutoff {
-                            self.source_marks.pop_front();
-                        } else {
-                            break;
-                        }
-                    }
-                    while let Some(&off) = self.disc_marks.front() {
-                        if off < cutoff {
-                            self.disc_marks.pop_front();
-                        } else {
-                            break;
-                        }
-                    }
-                }
-                break;
-            };
-            // The current AU ends where the next one begins: the first
-            // picture / sequence / GOP start code after this picture.
-            let end = match find_au_start(&self.buf, pic + 4) {
-                Some(b) => b,
-                None if force => self.buf.len(),
-                None if self.buf.len() > MAX_AU_BUFFER => self.buf.len(),
-                None => break, // AU not yet complete — await the next boundary
-            };
-            if end == 0 {
-                break;
-            }
-
-            // Phase 1 — read everything from `buf` before any mutation of self
-            // (the slice borrow must end before we touch self fields).
-            let hdr = extract_seq_header(&self.buf[..end]);
-            // A GOP header (0xB8) or a fresh sequence header (0xB3) starts a new
-            // GOP, resetting temporal_reference to 0.
-            let gop_boundary = find_code(&self.buf[..end], 0, GOP_CODE).is_some()
-                || find_code(&self.buf[..end], 0, SEQ_HEADER_CODE).is_some();
-            // picture_coding_type: the full 3-bit value (bits 5-3 of buf[pic+5]).
-            // 0 when the picture header is truncated (no coding type available).
-            let raw_coding_type = if pic + 5 < end {
-                (self.buf[pic + 5] >> 3) & 0x07
-            } else {
-                0
-            };
-            // temporal_reference: the 10 bits immediately after the picture
-            // start code = display order within the GOP.
-            let tr = if pic + 5 < end {
-                (((self.buf[pic + 4] as u64) << 2) | ((self.buf[pic + 5] as u64) >> 6)) & 0x3FF
-            } else {
-                0
-            };
-            let end_abs = self.base_offset + end as u64;
-            let data = self.buf[..end].to_vec();
-
-            // Phase 2 — mutate self.
-            if let Some(h) = hdr {
-                self.progressive_sequence = parse_progressive_sequence(&h);
-                self.seq_header = Some(h);
-                if let Some((num, den)) = self.frame_rate() {
-                    if num > 0 {
-                        self.frame_duration_ns = 1_000_000_000i64 * den as i64 / num as i64;
-                    }
-                }
-            }
-            // Decode the picture coding extension ONCE here and fold every
-            // per-picture datum (coding type + tff/rff/progressive_frame/
-            // frame_picture, plus the sequence's progressive flag) into one
-            // codec-agnostic `PictureInfo`. `nb_fields()`, `keyframe()`, and
-            // `field_order()` all derive from it; nothing downstream re-parses
-            // the elementary stream.
-            let (tff, rff, progressive_frame, frame_picture) = picture_coding_flags(&data);
-            let info = PictureInfo::mpeg2(
-                coding_type_from_raw(raw_coding_type),
-                Mpeg2Coding {
-                    top_field_first: tff,
-                    repeat_first_field: rff,
-                    progressive_frame,
-                    progressive_sequence: self.progressive_sequence,
-                    frame_picture,
-                },
-            );
-            let keyframe = info.keyframe();
-
-            // An explicit PES PTS for this access unit, if any. By the mark-drain
-            // invariant the front mark's offset is >= this AU's start, so a front
-            // mark inside [start, end) is this AU's own timestamp.
-            let explicit = self
-                .pts_marks
-                .front()
-                .filter(|&&(off, _)| off < end_abs)
-                .map(|&(_, p)| p);
-
-            // Byte-exact source provenance for this AU, by the same mark-drain
-            // invariant as the PTS: the front source mark inside [start, end)
-            // belongs to this access unit.
-            let src = self
-                .source_marks
-                .front()
-                .filter(|&&(off, _)| off < end_abs)
-                .map(|&(_, s)| s);
-
-            // A GOP boundary means the buffered run is a COMPLETE GOP (all its
-            // pictures display before the next GOP's), so flush it before
-            // starting the new one. `temporal_reference` resets to 0 at the
-            // boundary, keeping each GOP's display order self-contained.
-            if gop_boundary && !self.gop_buf.is_empty() {
-                self.flush_gop(&mut out);
-            }
-            // A concealed-gap mark inside this AU's range [start, end_abs) means
-            // this picture's own bytes begin after the gap — the first post-gap
-            // AU. Same front-mark invariant as PTS/source. Carries through GOP
-            // buffering/reorder to the ResyncGate (which arms at this picture).
-            let discontinuity = self.disc_marks.front().is_some_and(|&off| off < end_abs);
-            self.gop_buf.push(BufferedPicture {
-                tr,
-                info,
-                explicit_pts: explicit,
-                frame: Frame {
-                    pts_ns: 0,
-                    keyframe,
-                    discontinuity,
-                    data,
-                    duration_ns: None,
-                    coding: Some(info),
-                    source: src,
-                },
-            });
-            // Safety cap: a stream with no GOP/sequence boundaries would buffer
-            // unbounded. Force-flush a pathologically long run as its own GOP.
-            if self.gop_buf.len() >= MAX_PENDING_FRAMES {
-                self.flush_gop(&mut out);
-            }
-            self.buf.drain(..end);
-            self.base_offset = end_abs;
-            // Drop PTS marks fully consumed by the emitted AU; keep the mark at
-            // the boundary (it belongs to the next AU).
-            while let Some(&(off, _)) = self.pts_marks.front() {
-                if off < end_abs {
-                    self.pts_marks.pop_front();
-                } else {
-                    break;
-                }
-            }
-            while let Some(&(off, _)) = self.source_marks.front() {
-                if off < end_abs {
-                    self.source_marks.pop_front();
-                } else {
-                    break;
-                }
-            }
-            while let Some(&off) = self.disc_marks.front() {
-                if off < end_abs {
-                    self.disc_marks.pop_front();
-                } else {
-                    break;
                 }
             }
         }
-        // EOF: emit the final (possibly incomplete) GOP so nothing is dropped.
+        // A GOP header (0xB8) or a fresh sequence header (0xB3) starts a new GOP,
-        if force {
+        // resetting temporal_reference to 0.
-            self.flush_gop(&mut out);
+        let gop_boundary = find_code(&data, 0, GOP_CODE).is_some()
+            || find_code(&data, 0, SEQ_HEADER_CODE).is_some();
+        // picture_coding_type: the full 3-bit value (bits 5-3 of data[pic+5]).
+        // 0 when the picture header is truncated (no coding type available).
+        let raw_coding_type = if pic + 5 < end {
+            (data[pic + 5] >> 3) & 0x07
+        } else {
+            0
+        };
+        // temporal_reference: the 10 bits immediately after the picture start
+        // code = display order within the GOP.
+        let tr = if pic + 5 < end {
+            (((data[pic + 4] as u64) << 2) | ((data[pic + 5] as u64) >> 6)) & 0x3FF
+        } else {
+            0
+        };
+        // Decode the picture coding extension ONCE here and fold every
+        // per-picture datum (coding type + tff/rff/progressive_frame/
+        // frame_picture, plus the sequence's progressive flag) into one
+        // codec-agnostic `PictureInfo`. `nb_fields()`, `keyframe()`, and
+        // `field_order()` all derive from it; nothing downstream re-parses the
+        // elementary stream.
+        let (tff, rff, progressive_frame, frame_picture) = picture_coding_flags(&data);
+        let info = PictureInfo::mpeg2(
+            coding_type_from_raw(raw_coding_type),
+            Mpeg2Coding {
+                top_field_first: tff,
+                repeat_first_field: rff,
+                progressive_frame,
+                progressive_sequence: self.progressive_sequence,
+                frame_picture,
+            },
+        );
+        let keyframe = info.keyframe();
+
+        // A GOP boundary means the buffered run is a COMPLETE GOP (all its
+        // pictures display before the next GOP's), so flush it before starting
+        // the new one. `temporal_reference` resets to 0 at the boundary, keeping
+        // each GOP's display order self-contained.
+        if gop_boundary && !self.gop_buf.is_empty() {
+            self.flush_gop(out);
+        }
+        self.gop_buf.push(BufferedPicture {
+            tr,
+            info,
+            explicit_pts: au.pts,
+            frame: Frame {
+                pts_ns: 0,
+                keyframe,
+                // The assembler attributes the concealed-gap flag to the AU whose
+                // own bytes begin after the gap — the first post-gap picture — so
+                // it rides through GOP buffering/reorder to the ResyncGate.
+                discontinuity: au.discontinuity,
+                data,
+                duration_ns: None,
+                coding: Some(info),
+                source: au.source,
+            },
+        });
+        // Safety cap: a stream with no GOP/sequence boundaries would buffer
+        // unbounded. Force-flush a pathologically long run as its own GOP.
+        if self.gop_buf.len() >= MAX_PENDING_FRAMES {
+            self.flush_gop(out);
         }
-        out
     }
 
     /// Emit the buffered GOP. Each frame's PTS is the display-order prefix-sum of
@@ -448,31 +326,31 @@ impl CodecParser for Mpeg2Parser {
         if pes.data.is_empty() {
             return Vec::new();
         }
-        // Record this PES's timestamp against the absolute offset of its first
+        // Feed the fragment to the assembler, which reframes the elementary
-        // ES byte, BEFORE appending. MKV block timecodes are presentation
+        // stream on picture boundaries and hands back each complete access unit
-        // timestamps; prefer PTS (DTS shows B-frames in decode order — judder
+        // with its start timing. MKV block timecodes are presentation timestamps;
-        // and broken seeking), falling back to DTS only when PTS is absent.
+        // prefer PTS (DTS shows B-frames in decode order — judder and broken
-        let off = self.base_offset + self.buf.len() as u64;
+        // seeking), falling back to DTS only when PTS is absent.
-        if let Some(ts) = pes.pts.or(pes.dts) {
+        let pts = pes.pts.or(pes.dts).map(pts_to_ns);
-            self.pts_marks.push_back((off, pts_to_ns(ts)));
+        let aus = self
+            .au_asm
+            .push(&pes.data, pts, None, pes.source, pes.discontinuity);
+        let mut out = Vec::new();
+        for au in aus {
+            self.process_au(au, &mut out);
         }
-        if let Some(src) = pes.source {
+        out
-            self.source_marks.push_back((off, src));
-        }
-        // A concealed/lost gap on this PES marks the access unit its bytes begin —
-        // associated by offset (like PTS/source) so it lands on the first post-gap
-        // picture, not the previous one that completes when this PES arrives.
-        if pes.discontinuity {
-            self.disc_marks.push_back(off);
-        }
-        self.buf.extend_from_slice(&pes.data);
-        self.drain_complete_aus(false)
     }
 
     fn flush(&mut self) -> Vec<Frame> {
-        // drain_complete_aus(true) force-completes the trailing access unit and
+        // Force-complete the trailing access unit, then flush the final GOP so
-        // flushes the final GOP, so nothing is left buffered at EOF.
+        // nothing is left buffered at EOF.
-        self.drain_complete_aus(true)
+        let mut out = Vec::new();
+        for au in self.au_asm.flush() {
+            self.process_au(au, &mut out);
+        }
+        self.flush_gop(&mut out);
+        out
     }
 
     fn codec_private(&self) -> Option<Vec<u8>> {
@@ -518,25 +396,6 @@ fn find_code(data: &[u8], from: usize, want: u8) -> Option<usize> {
     None
 }
 
-/// Find the next access-unit boundary at or after `from`: the position of a
-/// picture (0x00), sequence header (0xB3), or GOP (0xB8) start code. Extension
-/// (0xB5), slice (0x01..=0xAF), user-data (0xB2) and sequence-end (0xB7) codes
-/// belong to the current access unit and are NOT boundaries.
-fn find_au_start(data: &[u8], from: usize) -> Option<usize> {
-    let mut pos = from;
-    while let Some(sc) = find_start_code(data, pos) {
-        if sc + 3 >= data.len() {
-            return None;
-        }
-        let code = data[sc + 3];
-        if code == PICTURE_CODE || code == SEQ_HEADER_CODE || code == GOP_CODE {
-            return Some(sc);
-        }
-        pos = sc + 4;
-    }
-    None
-}
-
 /// Parse horizontal and vertical resolution from sequence header bytes.
 /// The sequence header must start with 00 00 01 B3.
 fn parse_resolution(hdr: &[u8]) -> Option<(u16, u16)> {