From 3633882d6c99405f198a820964209d0377d26690 Mon Sep 17 00:00:00 2001 From: Matthew Jackson <1085847+MattJackson@users.noreply.github.com> Date: Thu, 9 Jul 2026 08:44:01 -0700 Subject: [PATCH] mux: reassemble MPEG-2 access units via the shared AuAssembler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 150 ++++++++++++++++- src/mux/codec/mpeg2.rs | 365 +++++++++++++---------------------------- 2 files changed, 255 insertions(+), 260 deletions(-) diff --git a/src/mux/au_assembly.rs b/src/mux/au_assembly.rs index f080613..14f2dbb 100644 --- a/src/mux/au_assembly.rs +++ b/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 `. +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 { 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 { // 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 { 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 { + 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 { + 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); diff --git a/src/mux/codec/mpeg2.rs b/src/mux/codec/mpeg2.rs index ae4f9c1..ba4bf8c 100644 --- a/src/mux/codec/mpeg2.rs +++ b/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>, - /// Unemitted elementary-stream bytes: the in-progress access unit plus any - /// lookahead needed to detect the next AU boundary. - buf: Vec, - /// Absolute ES byte offset of `buf[0]`. Used to associate PES PTS marks - /// (recorded by absolute offset) with the access units they belong to. - 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)>, + /// Reassembles PES fragments into complete access units (one coded picture + /// with its leading sequence/GOP headers) and carries each AU's start + /// timing / source / discontinuity forward — the shared machinery the + /// H.264/HEVC/VC-1 parsers also use, in its MPEG-2 mode. + au_asm: crate::mux::au_assembly::AuAssembler, /// 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, - /// 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, } /// 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), - base_offset: 0, - pts_marks: VecDeque::new(), - source_marks: VecDeque::new(), + au_asm: crate::mux::au_assembly::AuAssembler::mpeg2(), 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 { - let mut out = Vec::new(); - loop { - // An access unit must contain a coded picture; without one there is - // nothing to emit yet (leading sequence/GOP headers wait for it). - let Some(pic) = find_code(&self.buf, 0, PICTURE_CODE) else { - // No coded picture in an over-cap buffer means we are - // accumulating unparseable data (a stream with no picture - // start codes). Drop all but a 3-byte tail — enough to catch a - // start-code prefix straddling the boundary — and advance the - // absolute offset so the PES-mark invariant holds. Mirrors the - // post-picture buffer backstop in the AU-boundary search below. - if self.buf.len() > MAX_AU_BUFFER { - let drop = self.buf.len() - 3; - self.base_offset += drop as u64; - self.buf.drain(..drop); - let cutoff = self.base_offset; - while let Some(&(off, _)) = self.pts_marks.front() { - if off < cutoff { - self.pts_marks.pop_front(); - } 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; + /// Process one reassembled access unit (from [`AuAssembler`]): decode its + /// per-picture coding info, capture a new sequence header, and buffer the + /// picture into the current GOP for display-order timestamping. The AU's + /// timing / source / discontinuity were already attributed by the assembler. + fn process_au(&mut self, au: crate::mux::au_assembly::AssembledAu, out: &mut Vec) { + let data = au.data; + // An access unit must contain a coded picture; a fragment that assembled + // without one (only headers, or truncated at EOF) yields nothing. + let Some(pic) = find_code(&data, 0, PICTURE_CODE) else { + return; + }; + let end = data.len(); + // Capture a sequence header for codecPrivate; a new one replaces the + // stored value and re-locks the frame duration. + if let Some(h) = extract_seq_header(&data) { + 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; } } } - // EOF: emit the final (possibly incomplete) GOP so nothing is dropped. - if force { - self.flush_gop(&mut out); + // A GOP header (0xB8) or a fresh sequence header (0xB3) starts a new GOP, + // resetting temporal_reference to 0. + 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 - // ES byte, BEFORE appending. MKV block timecodes are presentation - // timestamps; prefer PTS (DTS shows B-frames in decode order — judder - // and broken seeking), falling back to DTS only when PTS is absent. - let off = self.base_offset + self.buf.len() as u64; - if let Some(ts) = pes.pts.or(pes.dts) { - self.pts_marks.push_back((off, pts_to_ns(ts))); + // Feed the fragment to the assembler, which reframes the elementary + // stream on picture boundaries and hands back each complete access unit + // with its start timing. MKV block timecodes are presentation timestamps; + // prefer PTS (DTS shows B-frames in decode order — judder and broken + // seeking), falling back to DTS only when PTS is absent. + let pts = pes.pts.or(pes.dts).map(pts_to_ns); + 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 { - 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) + out } fn flush(&mut self) -> Vec { - // drain_complete_aus(true) force-completes the trailing access unit and - // flushes the final GOP, so nothing is left buffered at EOF. - self.drain_complete_aus(true) + // Force-complete the trailing access unit, then flush the final GOP so + // nothing is left buffered at EOF. + 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> { @@ -518,25 +396,6 @@ fn find_code(data: &[u8], from: usize, want: u8) -> Option { 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 { - 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)> {