Fix DTS-HD MA muxing as lossy core only (#10)

DTS-HD MA/HRA access units on Blu-ray are a DTS core frame (sync
0x7FFE8001) followed by one or more DTS extension substreams (sync
0x64582025) carrying the lossless audio. Ground-truthing the Dunkirk
ISO showed the m2ts demuxer hands these out as SEPARATE PES packets on
the same PID: one core PES (exactly core-sized, nothing trailing), then
the extension substreams in following PES packets with their own later
PTS.

The old DtsParser emitted one frame per PES the moment a core frame was
complete, and dropped any PES with no core sync. So every core became a
core-only (lossy) frame and the extension PES packets were discarded as
junk -- silently downgrading the track to lossy DTS core (1557 kb/s CBR,
16-bit) instead of DTS-HD MA (VBR, 24-bit lossless).

Rewrite the parser to assemble across PES boundaries: an access unit
runs from its core sync up to (but not including) the NEXT core sync, so
the core plus every following extension substream stays together. Add a
CodecParser::flush() (default empty) called at end-of-stream by both the
pipelined and inline DiscStream mux paths to drain the final buffered
unit. A 64 KiB cap guarantees forward progress and never stalls if a
boundary can't be found.

Validated on the rip1 testbed: Dunkirk eng+ger and Fight Club eng main
audio now ffprobe as profile=DTS-HD MA (Fight Club eng at 24-bit), with
VBR packet sizes (~2716-2788 B) well above the old fixed 2012 B lossy
core. Genuinely-lossy DTS dub tracks are left untouched.

src/mux/codec/dts.rs

@@ -8,10 +8,20 @@
 use super::{CodecParser, Frame, PesPacket, pts_to_ns};
 
 const DTS_CORE_SYNC: [u8; 4] = [0x7F, 0xFE, 0x80, 0x01];
+/// DTS-HD extension substream syncword. The parser delimits an access unit by
+/// the next CORE sync (so every extension between two cores is captured), and
+/// never needs to locate or size the extension itself — so this is referenced
+/// only by the tests that synthesize extension substreams.
+#[cfg(test)]
 const DTS_HD_EXT_SYNC: [u8; 4] = [0x64, 0x58, 0x20, 0x25];
 
 pub struct DtsParser {
     buf: Vec<u8>,
+    /// PTS of the access unit currently being assembled in `buf` (the unit
+    /// starting at the first buffered core sync). Captured when that core
+    /// frame's PES first arrived; the trailing extension-substream PES
+    /// packets carry their own (later) PTS which must NOT override it.
+    pending_pts: i64,
 }
 
 impl Default for DtsParser {
@@ -24,10 +34,16 @@ impl DtsParser {
     pub fn new() -> Self {
         Self {
             buf: Vec::with_capacity(32768),
+            pending_pts: 0,
         }
     }
 }
 
+/// Hard cap on a buffered access unit (core + all its extension substreams).
+/// A DTS-HD MA frame is at most a few tens of KB; if the buffer grows past
+/// this without a clean boundary we resync rather than stall or balloon.
+const MAX_AU_BYTES: usize = 65536;
+
 impl CodecParser for DtsParser {
     fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
         if pes.data.is_empty() {
@@ -35,97 +51,120 @@ impl CodecParser for DtsParser {
         }
         let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0);
 
+        // On Blu-ray, a DTS-HD MA/HRA access unit is a DTS core frame
+        // (sync 0x7FFE8001) followed by one or more DTS extension substreams
+        // (sync 0x64582025). The m2ts demuxer hands those out as SEPARATE PES
+        // packets on the same PID — the core in one PES, then the extension
+        // substreams in following PES packets (with their own, later PTS). The
+        // lossless audio lives entirely in the extension substreams, so an
+        // access unit is only complete once all of its trailing extensions
+        // have been buffered. We assemble across PES boundaries here: an access
+        // unit runs from its core sync up to (but not including) the NEXT core
+        // sync. Emitting on the core boundary keeps the core + every following
+        // extension substream together (the lossless data), instead of the
+        // old per-PES emit that dropped the extension PES packets and
+        // downgraded the track to lossy DTS core (the Dunkirk / Fight Club
+        // bug). The PTS is the core frame's PTS, captured when the unit began.
+        if self.buf.is_empty() {
+            self.pending_pts = pts_ns;
+        }
         self.buf.extend_from_slice(&pes.data);
 
-        let data = &self.buf;
         let mut frames = Vec::new();
-        let mut pos = 0;
 
-        while pos < data.len() {
+        loop {
-            // Find DTS core sync
+            // Resync to the first core sync; drop any leading junk.
-            let start = match find_sync(&data[pos..], &DTS_CORE_SYNC) {
+            let Some(start) = find_sync(&self.buf, &DTS_CORE_SYNC) else {
-                Some(offset) => pos + offset,
+                // No core sync at all yet — keep at most a 3-byte tail so a
-                None => break,
+                // sync split across PES packets can still be found next time.
+                if self.buf.len() > 3 {
+                    let tail = self.buf.len() - 3;
+                    self.buf.drain(..tail);
+                }
+                break;
             };
-
+            if start > 0 {
-            // Need at least 10 bytes for core header to get frame size
+                self.buf.drain(..start);
-            if start + 10 > data.len() {
+                if find_sync(&self.buf, &DTS_CORE_SYNC) != Some(0) {
+                    // Shouldn't happen, but never loop forever.
                     break;
                 }
+            }
 
-            let core_size = dts_core_frame_size(&data[start..]);
+            // Need the core header to size the core frame.
-            if core_size == 0 || core_size > 32768 {
+            if self.buf.len() < 10 {
-                pos = start + 4;
+                break;
+            }
+            let core_size = dts_core_frame_size(&self.buf);
+            if core_size == 0 || core_size > MAX_AU_BYTES {
+                // Bogus core sync — skip past it and resync.
+                self.buf.drain(..4);
                 continue;
             }
+            if self.buf.len() < core_size {
+                break; // core frame not fully buffered yet — wait
+            }
 
-            if start + core_size > data.len() {
+            // The access unit ends at the next core sync. Search begins after
-                // Incomplete core frame
+            // this core's syncword so we don't re-match it. Anything between
+            // the core and that next sync is this unit's extension substream(s).
+            let au_end = match find_sync(&self.buf[core_size..], &DTS_CORE_SYNC) {
+                Some(rel) => core_size + rel,
+                None => {
+                    // No next core sync buffered yet. The trailing extension
+                    // substream PES packets may still be arriving, so WAIT for
+                    // them rather than emit a core-only (lossy) frame — unless
+                    // the buffer has grown unreasonably large, in which case
+                    // emit what we have to guarantee forward progress.
+                    if self.buf.len() <= MAX_AU_BYTES {
                         break;
                     }
-
+                    self.buf.len()
-            // Include the DTS-HD extension substream if one immediately follows
-            // the core. The extension carries the LOSSLESS (DTS-HD MA / HRA)
-            // data; emitting a core-only frame and dropping the trailing
-            // extension silently downgrades the track to lossy DTS core. If we
-            // can't yet tell whether an extension follows, or it's present but
-            // not fully buffered, WAIT for more PES data (break, leaving `pos`
-            // at this access unit's core sync so the buffer keeps the partial
-            // unit) rather than splitting the extension off and losing it.
-            let after = start + core_size;
-            let avail = data.len() - after;
-            // Does a DTS-HD extension substream follow the core? Match the full
-            // sync when it's buffered, or a partial PREFIX when the buffer ends
-            // mid-sync — so we wait for the rest instead of splitting the
-            // extension off and losing the lossless data. Nothing after the core
-            // (e.g. the final access unit / EOF, with no parser flush) is taken
-            // as a legitimate lossy core-only unit.
-            let ext_follows = if avail >= 4 {
-                data[after..after + 4] == DTS_HD_EXT_SYNC
-            } else if avail > 0 {
-                DTS_HD_EXT_SYNC[..avail] == data[after..]
-            } else {
-                false
-            };
-            let total_size = if !ext_follows {
-                core_size
-            } else if avail < 9 {
-                break; // extension present but its size header isn't buffered — wait
-            } else {
-                let ext_size = dts_hd_ext_frame_size(&data[after..]);
-                if ext_size == 0 || avail < ext_size {
-                    break; // extension known but not fully buffered — wait
                 }
-                core_size + ext_size
             };
 
+            let au: Vec<u8> = self.buf[..au_end].to_vec();
             frames.push(Frame {
-                pts_ns,
+                pts_ns: self.pending_pts,
                 keyframe: true,
-                data: data[start..start + total_size].to_vec(),
+                data: au,
                 duration_ns: None,
             });
-            pos = start + total_size;
+            self.buf.drain(..au_end);
-        }
+            // The next access unit (now at buf start) belongs to a later PTS.
-
+            // We can't know it exactly until its core PES arrives, but the
-        // Keep unconsumed data
+            // current PES's PTS is the best available approximation when the
-        let keep_from = if pos < data.len() {
+            // boundary fell inside this PES; refine on the next call's start.
-            find_sync(&data[pos..], &DTS_CORE_SYNC)
+            self.pending_pts = pts_ns;
-                .map(|o| pos + o)
-                .unwrap_or(data.len())
-        } else {
-            data.len()
-        };
-
-        if keep_from < data.len() {
-            self.buf = data[keep_from..].to_vec();
-        } else {
-            self.buf.clear();
         }
 
         frames
     }
 
+    fn flush(&mut self) -> Vec<Frame> {
+        // End of stream: emit the final access unit still buffered (the last
+        // core + its extension substreams, which had no following core sync to
+        // close it during streaming). Require a complete core frame; drop a
+        // bare partial sync tail.
+        if find_sync(&self.buf, &DTS_CORE_SYNC) != Some(0) || self.buf.len() < 10 {
+            self.buf.clear();
+            return Vec::new();
+        }
+        let core_size = dts_core_frame_size(&self.buf);
+        if core_size == 0 || self.buf.len() < core_size {
+            self.buf.clear();
+            return Vec::new();
+        }
+        let au = std::mem::take(&mut self.buf);
+        let pts_ns = self.pending_pts;
+        vec![Frame {
+            pts_ns,
+            keyframe: true,
+            data: au,
+            duration_ns: None,
+        }]
+    }
+
     fn codec_private(&self) -> Option<Vec<u8>> {
         None
     }
@@ -151,20 +190,6 @@ fn dts_core_frame_size(data: &[u8]) -> usize {
     fsize + 1
 }
 
-/// DTS-HD extension frame size from extension header.
-pub fn dts_hd_ext_frame_size(ext: &[u8]) -> usize {
-    if ext.len() < 9 {
-        return 0;
-    }
-    let raw =
-        ((ext[6] as usize & 0x1F) << 11) | ((ext[7] as usize) << 3) | ((ext[8] as usize) >> 5);
-    raw + 1
-}
-
-pub fn find_dts_hd_ext_sync(data: &[u8]) -> Option<usize> {
-    find_sync(data, &DTS_HD_EXT_SYNC)
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -198,12 +223,16 @@ mod tests {
 
     #[test]
     fn parse_single_frame() {
+        // A single core frame with no following core sync is the LAST access
+        // unit — held during streaming (can't know an extension won't follow),
+        // then drained on flush() at EOF.
         let mut parser = DtsParser::new();
         let frame = make_dts_core(512);
         let pes = make_pes(frame, Some(90000));
-        let frames = parser.parse(&pes);
+        assert!(parser.parse(&pes).is_empty());
-        assert_eq!(frames.len(), 1);
+        let tail = parser.flush();
-        assert_eq!(frames[0].data.len(), 512);
+        assert_eq!(tail.len(), 1);
+        assert_eq!(tail[0].data.len(), 512);
     }
 
     #[test]
@@ -216,50 +245,108 @@ mod tests {
         assert!(parser.parse(&pes1).is_empty());
 
         let pes2 = make_pes(frame[mid..].to_vec(), Some(93000));
-        let frames = parser.parse(&pes2);
+        assert!(parser.parse(&pes2).is_empty());
-        assert_eq!(frames.len(), 1);
+        let tail = parser.flush();
-        assert_eq!(frames[0].data.len(), 512);
+        assert_eq!(tail.len(), 1);
+        assert_eq!(tail[0].data.len(), 512);
     }
 
-    /// Build a DTS-HD extension substream of `size` bytes with a valid sync +
+    #[test]
-    /// size header (matching `dts_hd_ext_frame_size`).
+    fn two_cores_back_to_back_emit_first_on_boundary() {
+        // The first complete unit is emitted as soon as the next core sync is
+        // seen; the second is held until flush.
+        let mut parser = DtsParser::new();
+        let mut stream = make_dts_core(512);
+        stream.extend_from_slice(&make_dts_core(640));
+        let f = parser.parse(&make_pes(stream, Some(90000)));
+        assert_eq!(f.len(), 1);
+        assert_eq!(f[0].data.len(), 512);
+        let tail = parser.flush();
+        assert_eq!(tail.len(), 1);
+        assert_eq!(tail[0].data.len(), 640);
+    }
+
+    /// Build a minimal DTS-HD extension substream of `size` bytes (just the
+    /// sync + zero-padding). The parser delimits extensions by the next CORE
+    /// sync, not by the extension's own size header, so a valid header isn't
+    /// required — only that the bytes carry no spurious core sync.
     fn make_dts_ext(size: usize) -> Vec<u8> {
-        let raw = size - 1;
         let mut e = vec![0u8; size];
         e[0..4].copy_from_slice(&DTS_HD_EXT_SYNC);
-        e[6] = ((raw >> 11) & 0x1F) as u8;
-        e[7] = ((raw >> 3) & 0xFF) as u8;
-        e[8] = (((raw & 0x07) << 5) as u8) | (e[8] & 0x1F);
         e
     }
 
     #[test]
-    fn keeps_dts_hd_extension_across_pes_boundary() {
+    fn keeps_dts_hd_extension_in_separate_pes_packets() {
-        // DTS-HD MA access unit = core + extension substream. When the
+        // The real Blu-ray layout (ground-truthed on Dunkirk): the DTS core
-        // extension straddles a PES boundary, the parser must WAIT and emit the
+        // arrives in one PES, then its DTS-HD MA extension substreams arrive
-        // full unit — not a core-only frame (which would drop the lossless
+        // in SEPARATE following PES packets on the same PID. The parser must
-        // data, the Dunkirk lossy-core bug).
+        // stitch core + all trailing extensions into one access unit — not
+        // emit a core-only (lossy) frame and drop the extension PES packets
+        // (the Dunkirk / Fight Club lossy-core bug).
         let mut parser = DtsParser::new();
-        let core = make_dts_core(512);
-        let ext = make_dts_ext(256);
-        let mut au = core;
-        au.extend_from_slice(&ext); // 768-byte access unit
 
-        // Split mid-extension: first PES carries the core + 100 ext bytes.
+        // Frame 1: core (512) + two extension substreams (256 + 200).
-        let split = 512 + 100;
-        let f1 = parser.parse(&make_pes(au[..split].to_vec(), Some(90000)));
         assert!(
-            f1.is_empty(),
+            parser
-            "must wait for the full extension, not emit a core-only frame"
+                .parse(&make_pes(make_dts_core(512), Some(90000)))
+                .is_empty(),
+            "core alone: must wait for any following extension"
+        );
+        assert!(
+            parser
+                .parse(&make_pes(make_dts_ext(256), Some(91000)))
+                .is_empty(),
+            "first extension PES: still waiting for the unit to close"
+        );
+        assert!(
+            parser
+                .parse(&make_pes(make_dts_ext(200), Some(91500)))
+                .is_empty(),
+            "second extension PES: unit still not closed (no next core yet)"
         );
 
-        let f2 = parser.parse(&make_pes(au[split..].to_vec(), Some(90000)));
+        // Frame 2's core PES arrives — that closes frame 1. The emitted unit
-        assert_eq!(f2.len(), 1);
+        // must be core + BOTH extensions (lossless preserved), and keep the
+        // core's PTS, not the extension PES timestamps.
+        let f = parser.parse(&make_pes(make_dts_core(512), Some(93000)));
+        assert_eq!(f.len(), 1);
         assert_eq!(
-            f2[0].data.len(),
+            f[0].data.len(),
-            768,
+            512 + 256 + 200,
-            "frame must include core + extension (lossless preserved)"
+            "frame must include core + every extension substream"
         );
+
+        // EOF drains frame 2.
+        let tail = parser.flush();
+        assert_eq!(tail.len(), 1);
+        assert_eq!(tail[0].data.len(), 512);
+    }
+
+    #[test]
+    fn extension_split_across_pes_is_preserved() {
+        // An extension substream straddling a PES boundary must still be fully
+        // attached to its core.
+        let mut parser = DtsParser::new();
+        let ext = make_dts_ext(300);
+        assert!(
+            parser
+                .parse(&make_pes(make_dts_core(512), Some(90000)))
+                .is_empty()
+        );
+        assert!(
+            parser
+                .parse(&make_pes(ext[..150].to_vec(), Some(91000)))
+                .is_empty()
+        );
+        assert!(
+            parser
+                .parse(&make_pes(ext[150..].to_vec(), Some(91000)))
+                .is_empty()
+        );
+        let tail = parser.flush();
+        assert_eq!(tail.len(), 1);
+        assert_eq!(tail[0].data.len(), 512 + 300);
     }
 
     #[test]

src/mux/codec/mod.rs

@@ -52,6 +52,18 @@ pub trait CodecParser: Send {
     /// Some (TrueHD): multiple access units per PES.
     fn parse(&mut self, pes: &PesPacket) -> Vec<Frame>;
 
+    /// Drain any access unit still buffered after the last PES.
+    ///
+    /// Parsers that buffer across PES boundaries to assemble a complete
+    /// access unit (e.g. DTS-HD, whose extension substreams arrive in
+    /// separate PES packets) hold the final unit until they can prove it's
+    /// complete. At end-of-stream there is no following packet to prove it,
+    /// so the demuxer calls `flush()` once after the last PES to emit it.
+    /// Default: nothing buffered, no tail.
+    fn flush(&mut self) -> Vec<Frame> {
+        Vec::new()
+    }
+
     /// Get codec initialization data (e.g., SPS+PPS for H.264).
     /// Returns None until enough data has been seen.
     fn codec_private(&self) -> Option<Vec<u8>>;

src/mux/disc.rs

@@ -541,6 +541,19 @@ impl crate::pes::Stream for DiscStream {
                         }
                     }
                 }
+                // Drain any access unit a codec parser buffered past the last
+                // PES (DTS-HD's final core+extension unit, assembled across
+                // PES boundaries).
+                let pid_to_track = &self.pid_to_track;
+                let pending = &mut self.pending_frames;
+                for (pid, parser) in self.parsers.iter_mut() {
+                    let Some(&(_, track)) = pid_to_track.iter().find(|(p, _)| p == pid) else {
+                        continue;
+                    };
+                    for frame in parser.flush() {
+                        pending.push_back(crate::pes::PesFrame::from_codec_frame(track, frame));
+                    }
+                }
                 return Ok(self.pending_frames.pop_front());
             }
 

src/mux/pipelined_stream.rs

@@ -176,6 +176,18 @@ impl Stream for PipelinedPesStream {
                 }
                 false => {
                     self.eof = true;
+                    // Drain any access unit a parser buffered past the last
+                    // PES (e.g. DTS-HD's final core+extension unit).
+                    let pid_to_track = &self.pid_to_track;
+                    let pending = &mut self.pending_frames;
+                    for (pid, parser) in self.parsers.iter_mut() {
+                        let Some(&(_, track)) = pid_to_track.iter().find(|(p, _)| p == pid) else {
+                            continue;
+                        };
+                        for frame in parser.flush() {
+                            pending.push_back(PesFrame::from_codec_frame(track, frame));
+                        }
+                    }
                     return Ok(self.pending_frames.pop_front());
                 }
             }