pq/libfreemkv
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mux: resume AU-boundary scans from a cursor (O(n) not O(n²))

Browse Source 
AuAssembler::drain rescanned the whole buffered access unit from a fixed
offset on every push, so reassembling one AU split across N program-stream
PES fragments cost O(bytes²/fragment) — amplified on the HD-DVD PS path
where H.264/HEVC/VC-1 frames are large and now flow through this shared
assembler (unlike the TS path, which delivers one AU per PES).

Carry a scan_pos cursor (and, for the stateful VC-1/MPEG-2 rules, a
seen_unit flag) so each push resumes the boundary search where the last
one stopped instead of restarting. Total scan work for one AU is now
O(AU bytes). The from-scratch scanners are retained as a #[cfg(test)]
oracle; a new differential test asserts the resumable path yields
byte-identical AUs at every fragment granularity for all three modes.
This commit is contained in:
Matthew Jackson
2026-07-09 16:22:54 -07:00
parent 26423187d3
commit b5a5138569
1 changed files with 155 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mux/au_assembly.rs
+155 -79
View File
@@ -102,6 +102,18 @@ pub(crate) struct AuAssembler {
marks: VecDeque<Mark>,
/// Absolute offsets of fragments flagged with an upstream discontinuity.
disc_marks: VecDeque<u64>,
/// Incremental boundary-scan cursor: the offset into `buf` up to which the
/// current AU has already been searched for its end without finding one. Each
/// `push` resumes the boundary search from here instead of rescanning the
/// whole buffer, so reassembling one AU split across N PES fragments costs
/// O(AU bytes) total, not O(AU bytes²/fragment). Reset to 0 whenever `buf[0]`
/// moves (an AU drained, or leading bytes dropped).
scan_pos: usize,
/// Whether the current AU has already contained a coded frame/picture — the
/// state the VC-1/MPEG-2 boundary rule carries across a resumed scan (their
/// boundary is "the next opener after a frame is already seen"). Meaningless
/// for `Mode::StartCode`. Reset with `scan_pos`.
seen_unit: bool,
}
impl AuAssembler {
@@ -122,6 +134,8 @@ impl AuAssembler {
base: 0,
marks: VecDeque::new(),
disc_marks: VecDeque::new(),
scan_pos: 0,
seen_unit: false,
}
}
@@ -136,6 +150,8 @@ impl AuAssembler {
base: 0,
marks: VecDeque::new(),
disc_marks: VecDeque::new(),
scan_pos: 0,
seen_unit: false,
}
}
@@ -199,6 +215,7 @@ impl AuAssembler {
let drop = self.buf.len() - 3;
self.buf.drain(..drop);
self.base += drop as u64;
self.reset_scan();
self.drop_marks_before(self.base);
}
break;
@@ -208,11 +225,14 @@ impl AuAssembler {
// before we synced (or junk) — discard them and any stale marks.
self.buf.drain(..a0);
self.base += a0 as u64;
self.reset_scan();
self.drop_marks_before(self.base);
continue;
}
// The AU runs from here (buf[0]) to the NEXT AU boundary.
let end = match au_boundary(mode, &self.buf) {
// The AU runs from here (buf[0]) to the NEXT AU boundary. The search
// resumes from `scan_pos` (bytes already searched with no boundary),
// so one AU spread across many fragments is scanned once, not per push.
let end = match self.au_boundary_resumable() {
Some(next) => next,
// No next boundary yet: on EOF (or over-cap backstop) the rest of
// the buffer is this AU; otherwise wait for more data.
@@ -250,6 +270,7 @@ impl AuAssembler {
let data = self.buf[..end].to_vec();
self.buf.drain(..end);
self.base += end as u64;
self.reset_scan();
out.push(AssembledAu {
data,
pts,
@@ -261,6 +282,72 @@ impl AuAssembler {
out
}
/// Reset the incremental boundary-scan cursor. Called whenever `buf[0]` moves
/// (an AU drained, or leading bytes discarded) so the next scan starts fresh
/// from the new AU opener.
fn reset_scan(&mut self) {
self.scan_pos = 0;
self.seen_unit = false;
}
/// Find the end of the AU that opens at `buf[0]`, resuming from `scan_pos`
/// (and, for VC-1/MPEG-2, the carried `seen_unit`) instead of rescanning the
/// whole buffer. On no boundary yet, advances `scan_pos`/`seen_unit` so the
/// next call continues where this one stopped. Equivalent result to a
/// from-scratch whole-buffer scan, but O(total AU bytes) across all pushes.
fn au_boundary_resumable(&mut self) -> Option<usize> {
match self.mode {
Mode::StartCode(marker) => {
// Stateless: the AU ends at the next delimiter after the opener at
// buf[0]. Resume from the furthest searched offset (never before 4,
// to skip the opening delimiter). find_start_code needs 4 bytes, so
// back up 3 to catch a code straddling the previous buffer end.
let from = self.scan_pos.max(4);
match find_start_code(&self.buf, from, marker) {
Some(e) => Some(e),
None => {
self.scan_pos = self.buf.len().saturating_sub(3).max(from);
None
}
}
}
Mode::Vc1 => self.scan_unit_boundary(VC1_FRAME, &[VC1_ENTRY, VC1_SEQ]),
Mode::Mpeg2 => self.scan_unit_boundary(MP2_PICTURE, &[MP2_SEQ, MP2_GOP]),
Mode::Passthrough => None,
}
}
/// Resumable form of the VC-1/MPEG-2 boundary rule: scan from `scan_pos`,
/// carrying `seen_unit`; the AU ends at the next `frame` / `header` start code
/// once a frame is already seen. Advances `scan_pos`/`seen_unit` when no
/// boundary is found so the next push continues, not restarts.
fn scan_unit_boundary(&mut self, frame: u8, headers: &[u8]) -> Option<usize> {
let buf = &self.buf;
let mut i = self.scan_pos;
let mut seen = self.seen_unit;
while i + 4 <= buf.len() {
if buf[i] == 0 && buf[i + 1] == 0 && buf[i + 2] == 1 {
let c = buf[i + 3];
let is_frame = c == frame;
if (is_frame || headers.contains(&c)) && i > 0 && seen {
// The AU ends at the next frame/header once a frame is seen.
return Some(i);
}
if is_frame {
seen = true;
}
i += 4;
} else {
i += 1;
}
}
// No boundary yet. Persist the scan state so the next append resumes here
// rather than rescanning from 0 (the i+=4 stride is preserved exactly).
self.scan_pos = i;
self.seen_unit = seen;
None
}
fn drop_marks_before(&mut self, off: u64) {
while self.marks.front().is_some_and(|m| m.off < off) {
self.marks.pop_front();
@@ -284,18 +371,6 @@ fn au_opener(mode: Mode, buf: &[u8]) -> Option<usize> {
}
}
/// Offset where the AU that opens at `buf[0]` ends (the start of the next AU), or
/// `None` if the next boundary is not yet buffered.
fn au_boundary(mode: Mode, buf: &[u8]) -> Option<usize> {
match mode {
// 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,
}
}
/// Find the next `00 00 01 <marker>` start code at or after `from`.
fn find_start_code(buf: &[u8], from: usize, marker: u8) -> Option<usize> {
let mut i = from;
@@ -325,38 +400,6 @@ fn find_vc1_start(buf: &[u8], from: usize) -> Option<usize> {
None
}
/// End offset of the VC-1 access unit that opens at `buf[0]`: the next
/// sequence-header / entry-point / frame BDU that appears *after* this AU already
/// contains a frame (`0x0D`). Returns `None` while the AU is still open (no frame
/// yet, or no following BDU buffered). A leading `0x0F`/`0x0E` header group thus
/// stays attached to the frame it precedes rather than the previous AU.
fn find_vc1_au_end(buf: &[u8]) -> Option<usize> {
let mut seen_frame = 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] {
VC1_FRAME => {
if i > 0 && seen_frame {
return Some(i);
}
seen_frame = true;
}
VC1_ENTRY | VC1_SEQ => {
if i > 0 && seen_frame {
return Some(i);
}
}
_ => {}
}
i += 4;
} else {
i += 1;
}
}
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> {
@@ -374,39 +417,6 @@ fn find_mpeg2_start(buf: &[u8], from: usize) -> Option<usize> {
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::*;
@@ -637,6 +647,72 @@ mod tests {
assert_eq!(out[0].data, full);
}
/// Split `stream` into fragments of `frag` bytes, push them through the given
/// assembler mode, and return the reassembled AU byte-payloads.
fn reassemble_with(mut a: AuAssembler, stream: &[u8], frag: usize) -> Vec<Vec<u8>> {
let mut out = Vec::new();
let mut i = 0;
while i < stream.len() {
let end = (i + frag).min(stream.len());
for au in a.push(&stream[i..end], None, None, None, false) {
out.push(au.data);
}
i = end;
}
for au in a.flush() {
out.push(au.data);
}
out
}
#[test]
fn resumable_boundary_matches_from_scratch_across_all_fragmentations() {
// The incremental scan_pos cursor must produce byte-identical AUs to a
// whole-buffer rescan, at EVERY fragment granularity (this is what makes
// the O(n) resume equivalent to the old O(n^2) from-scratch scan). Build a
// multi-AU stream per codec, reassemble it fed 1 byte at a time up to
// whole, and require one canonical result.
let h264 = {
let mut s = au(0x11, 40); // AU1 (AUD + payload)
s.extend(au(0x22, 70)); // AU2
s.extend(au(0x33, 25)); // AU3
s
};
let vc1 = {
let mut s = bdu(VC1_SEQ, 0xAA, 8);
s.extend(bdu(VC1_ENTRY, 0xBB, 6));
s.extend(bdu(VC1_FRAME, 0xCC, 50)); // I-frame AU
s.extend(bdu(VC1_FRAME, 0xDD, 30)); // P-frame AU
s.extend(bdu(VC1_FRAME, 0xEE, 20)); // P-frame AU
s
};
let mpeg2 = {
let mut s = bdu(MP2_SEQ, 0xAA, 10);
s.extend(bdu(MP2_GOP, 0xBB, 8));
s.extend(bdu(MP2_PICTURE, 0xCC, 60)); // GOP-opening picture AU
s.extend(bdu(MP2_PICTURE, 0xDD, 40)); // picture AU
s
};
// (label, stream, assembler factory). MPEG-2 uses the dedicated mpeg2()
// assembler (Mode::Mpeg2); the AUD/VC-1 codecs use for_codec().
let cases: [(&str, &[u8], fn() -> AuAssembler); 3] = [
("h264", &h264, || AuAssembler::for_codec(Codec::H264)),
("vc1", &vc1, || AuAssembler::for_codec(Codec::Vc1)),
("mpeg2", &mpeg2, AuAssembler::mpeg2),
];
for (label, stream, make) in cases {
let whole = reassemble_with(make(), stream, stream.len());
assert!(!whole.is_empty(), "{label}: baseline produced AUs");
for frag in 1..=stream.len() {
let got = reassemble_with(make(), stream, frag);
assert_eq!(
got, whole,
"{label}: fragmented at {frag} differs from whole-buffer reassembly"
);
}
}
}
#[test]
fn over_cap_without_boundary_force_flushes() {
let mut a = AuAssembler::for_codec(Codec::H264);