diff --git a/src/mux/ebml.rs b/src/mux/ebml.rs index 811c660..c79382b 100644 --- a/src/mux/ebml.rs +++ b/src/mux/ebml.rs @@ -184,6 +184,55 @@ pub fn end_master(w: &mut W, size_pos: u64) -> io::Result<()> { Ok(()) } +/// Start a master element **in an in-memory buffer**: append ID + the same +/// 8-byte size placeholder [`start_master`] writes. Returns the buffer index of +/// the size field, for [`end_master_buf`]. +/// +/// This is the seek-free twin of [`start_master`]/[`end_master`] for elements +/// small enough to assemble whole before hitting the file (a BlockGroup: one +/// Block plus a couple of tiny elements). Because [`end_master`] always +/// back-patches a FIXED-WIDTH 8-byte VINT (`0x01` + 7 payload bytes) rather +/// than a minimal-width one, the bytes produced by this pair are byte-for-byte +/// identical to the seek-and-back-patch pair — assembling in memory cannot +/// change the emitted Matroska. +pub fn start_master_buf(buf: &mut Vec, id: u32) -> io::Result { + write_id(buf, id)?; + let size_pos = buf.len(); + // 8-byte size placeholder (overwritten by end_master_buf) + buf.extend_from_slice(&[0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]); + Ok(size_pos) +} + +/// End a master element in an in-memory buffer: patch the 8-byte size field at +/// `size_pos` (as returned by [`start_master_buf`]) with the body length. +/// +/// Errors rather than panicking if `size_pos` does not name a placeholder that +/// is still inside the buffer, or if the body exceeds the 7-byte VINT payload. +pub fn end_master_buf(buf: &mut [u8], size_pos: usize) -> io::Result<()> { + let end = buf.len(); + let Some(body_start) = size_pos.checked_add(8) else { + return Err(crate::error::Error::MkvInvalid.into()); + }; + if end < body_start { + return Err(crate::error::Error::MkvInvalid.into()); + } + let data_size = (end - body_start) as u64; + if data_size >= 0x0100_0000_0000_0000 { + return Err(crate::error::Error::MkvInvalid.into()); + } + buf[size_pos..body_start].copy_from_slice(&[ + 0x01, + (data_size >> 48) as u8, + (data_size >> 40) as u8, + (data_size >> 32) as u8, + (data_size >> 24) as u8, + (data_size >> 16) as u8, + (data_size >> 8) as u8, + data_size as u8, + ]); + Ok(()) +} + // ============================================================ // EBML Read primitives // ============================================================ @@ -1330,4 +1379,90 @@ mod tests { // And the whole buffer is exactly the outer element. assert_eq!(data.len() as u64, outer_body_start + osize); } + /// The buffer-based master helpers must produce byte-for-byte what the + /// seek-based ones produce. `write_block_group` was rewritten onto + /// `start_master_buf`/`end_master_buf` to eliminate ~4 seeks (and 4 MiB + /// BufWriter flushes) PER FRAME, which on an MPEG-2 title is ~350k frames — a + /// change only safe because the two encodings are identical. + /// + /// Both write a FIXED-width 8-byte VINT placeholder (0x01 + 7 payload bytes) + /// and patch it in place; neither ever emits a minimal-width size. This test + /// pins that, so a future "optimisation" of either one to minimal-width sizes + /// cannot silently desync the two and change emitted Matroska. + #[test] + fn buffered_master_matches_seeking_master_byte_for_byte() { + use std::io::Cursor; + + // Bodies chosen to cross VINT-relevant magnitudes: empty, tiny, and one + // spanning more than a byte of length. + for body in [ + Vec::new(), + vec![0xAAu8], + (0..300u32).map(|i| (i % 251) as u8).collect::>(), + ] { + // Seek-based path. + let mut c = Cursor::new(Vec::new()); + let pos = start_master(&mut c, BLOCK_GROUP).unwrap(); + c.write_all(&body).unwrap(); + end_master(&mut c, pos).unwrap(); + let seeking = c.into_inner(); + + // Buffer-based path. + let mut buf = Vec::new(); + let bpos = start_master_buf(&mut buf, BLOCK_GROUP).unwrap(); + buf.extend_from_slice(&body); + end_master_buf(&mut buf, bpos).unwrap(); + + assert_eq!( + buf, + seeking, + "buffered and seeking master encodings diverged for a {}-byte body", + body.len() + ); + } + } + + /// Nested masters must patch correctly too — the MVC BlockGroup nests + /// BlockAdditions > BlockMore inside the BlockGroup, and in-memory patching + /// works by index rather than by file offset, so nesting is where an + /// index-arithmetic slip would show up. + #[test] + fn buffered_master_nests_correctly() { + use std::io::Cursor; + + let mut c = Cursor::new(Vec::new()); + let outer = start_master(&mut c, BLOCK_GROUP).unwrap(); + c.write_all(&[0x11, 0x22]).unwrap(); + let inner = start_master(&mut c, BLOCK_ADDITIONS).unwrap(); + c.write_all(&[0x33, 0x44, 0x55]).unwrap(); + end_master(&mut c, inner).unwrap(); + c.write_all(&[0x66]).unwrap(); + end_master(&mut c, outer).unwrap(); + let seeking = c.into_inner(); + + let mut buf = Vec::new(); + let outer_b = start_master_buf(&mut buf, BLOCK_GROUP).unwrap(); + buf.extend_from_slice(&[0x11, 0x22]); + let inner_b = start_master_buf(&mut buf, BLOCK_ADDITIONS).unwrap(); + buf.extend_from_slice(&[0x33, 0x44, 0x55]); + end_master_buf(&mut buf, inner_b).unwrap(); + buf.extend_from_slice(&[0x66]); + end_master_buf(&mut buf, outer_b).unwrap(); + + assert_eq!( + buf, seeking, + "nested buffered masters must match the seeking form" + ); + } + + /// `end_master_buf` must REFUSE a bogus position rather than panic on the + /// slice index — this crate must not panic from library code. + #[test] + fn end_master_buf_rejects_a_position_outside_the_buffer() { + let mut buf = vec![0u8; 4]; + assert!( + end_master_buf(&mut buf, 99).is_err(), + "a size_pos past the end of the buffer must error, not panic" + ); + } } diff --git a/src/mux/mkv.rs b/src/mux/mkv.rs index b885735..22d71fa 100644 --- a/src/mux/mkv.rs +++ b/src/mux/mkv.rs @@ -617,6 +617,20 @@ pub struct MkvMuxer { cluster_pos: u64, cluster_size_pos: u64, cluster_ts_ticks: i64, + /// Reusable scratch buffer for assembling ONE BlockGroup before it is + /// written with a single `write_all`. + /// + /// The BlockGroup path is the hot one — the MPEG-2 parser stamps a per-frame + /// duration, so every I/P/B frame lands there, plus AC-3 audio and PGS + /// subtitles (~350k BlockGroups for a 90-minute feature). Building the + /// element in memory means its size is known before it reaches the file, so + /// no `start_master`/`end_master` back-patch is needed: that removes the + /// per-frame `stream_position()` + two `seek()`s, each of which flushed the + /// 4 MiB `BufWriter` (`BufWriter` does not override `stream_position`, so a + /// position query is `seek(Current(0))` = flush + lseek) and reset the + /// writeback pipeline's `last_flush_pos`. Kept on the muxer so the + /// allocation is made once, not per frame. + block_group_buf: Vec, base_pts_ticks: Option, /// Last block timecode (TimestampScale ticks, relative to base_pts) written /// PER TRACK, to enforce strictly-monotonic per-track timestamps — @@ -1237,6 +1251,7 @@ impl MkvMuxer { cluster_open: false, cluster_pos: 0, cluster_size_pos: 0, + block_group_buf: Vec::new(), cluster_ts_ticks: 0, base_pts_ticks: None, last_pts_ticks: std::collections::HashMap::new(), @@ -1851,24 +1866,69 @@ impl MkvMuxer { data: &[u8], reference: Option, duration_ticks: u64, + ) -> io::Result<()> { + let mut buf = std::mem::take(&mut self.block_group_buf); + buf.clear(); + let res = Self::build_block_group( + &mut buf, + track_num, + relative_ts, + data, + reference, + Some(duration_ticks), + None, + ) + .and_then(|()| self.writer.write_all(&buf)); + // Hand the (now grown) scratch buffer back so the next frame reuses the + // allocation, on the error path too. + self.block_group_buf = buf; + res + } + + /// Assemble a complete BlockGroup into `buf` — one `write_all` at the call + /// site, no `stream_position` and no `seek`. + /// + /// `additional` is `Some(dependent AU)` for the MVC path, which appends + /// `BlockAdditions > BlockMore { BlockAddID=2, BlockAdditional }`. + fn build_block_group( + buf: &mut Vec, + track_num: usize, + relative_ts: i16, + data: &[u8], + reference: Option, + duration_ticks: Option, + additional: Option<&[u8]>, ) -> io::Result<()> { let (tv, tv_len) = track_vint(track_num); let track_vint = &tv[..tv_len]; + // The 0x80 Keyframe flag is SimpleBlock-only; inside a BlockGroup Block + // it is reserved and MUST be 0 — keyframe-ness is signalled by the + // presence/absence of ReferenceBlock. let flags: u8 = 0x00; let block_size = track_vint.len() + 2 + 1 + data.len(); - let bg_pos = ebml::start_master(&mut self.writer, ebml::BLOCK_GROUP)?; - ebml::write_id(&mut self.writer, ebml::BLOCK)?; - ebml::write_size(&mut self.writer, block_size as u64)?; - self.writer.write_all(track_vint)?; - self.writer.write_all(&relative_ts.to_be_bytes())?; - self.writer.write_all(&[flags])?; - self.writer.write_all(data)?; - ebml::write_uint(&mut self.writer, ebml::BLOCK_DURATION, duration_ticks)?; - if let Some(ref_off) = reference { - ebml::write_int(&mut self.writer, ebml::REFERENCE_BLOCK, ref_off)?; + let bg_pos = ebml::start_master_buf(buf, ebml::BLOCK_GROUP)?; + ebml::write_id(buf, ebml::BLOCK)?; + ebml::write_size(buf, block_size as u64)?; + buf.extend_from_slice(track_vint); + buf.extend_from_slice(&relative_ts.to_be_bytes()); + buf.push(flags); + buf.extend_from_slice(data); + if let Some(dt) = duration_ticks { + ebml::write_uint(buf, ebml::BLOCK_DURATION, dt)?; } - ebml::end_master(&mut self.writer, bg_pos)?; + if let Some(ref_off) = reference { + ebml::write_int(buf, ebml::REFERENCE_BLOCK, ref_off)?; + } + if let Some(additional) = additional { + let adds_pos = ebml::start_master_buf(buf, ebml::BLOCK_ADDITIONS)?; + let more_pos = ebml::start_master_buf(buf, ebml::BLOCK_MORE)?; + ebml::write_uint(buf, ebml::BLOCK_ADD_ID, BLOCK_ADD_ID_VALUE_MVC)?; + ebml::write_binary(buf, ebml::BLOCK_ADDITIONAL, additional)?; + ebml::end_master_buf(buf, more_pos)?; + ebml::end_master_buf(buf, adds_pos)?; + } + ebml::end_master_buf(buf, bg_pos)?; Ok(()) } @@ -1886,36 +1946,20 @@ impl MkvMuxer { reference: Option, duration_ticks: Option, ) -> io::Result<()> { - let (tv, tv_len) = track_vint(track_num); - let track_vint = &tv[..tv_len]; - // The 0x80 Keyframe flag is SimpleBlock-only; inside a BlockGroup Block - // it is reserved and MUST be 0 — keyframe-ness is signalled by the - // presence/absence of ReferenceBlock. - let flags: u8 = 0x00; - let block_size = track_vint.len() + 2 + 1 + data.len(); - - let bg_pos = ebml::start_master(&mut self.writer, ebml::BLOCK_GROUP)?; - ebml::write_id(&mut self.writer, ebml::BLOCK)?; - ebml::write_size(&mut self.writer, block_size as u64)?; - self.writer.write_all(track_vint)?; - self.writer.write_all(&relative_ts.to_be_bytes())?; - self.writer.write_all(&[flags])?; - self.writer.write_all(data)?; - if let Some(dt) = duration_ticks { - ebml::write_uint(&mut self.writer, ebml::BLOCK_DURATION, dt)?; - } - if let Some(ref_off) = reference { - ebml::write_int(&mut self.writer, ebml::REFERENCE_BLOCK, ref_off)?; - } - // BlockAdditions → BlockMore { BlockAddID=2, BlockAdditional=dependent AU }. - let adds_pos = ebml::start_master(&mut self.writer, ebml::BLOCK_ADDITIONS)?; - let more_pos = ebml::start_master(&mut self.writer, ebml::BLOCK_MORE)?; - ebml::write_uint(&mut self.writer, ebml::BLOCK_ADD_ID, BLOCK_ADD_ID_VALUE_MVC)?; - ebml::write_binary(&mut self.writer, ebml::BLOCK_ADDITIONAL, additional)?; - ebml::end_master(&mut self.writer, more_pos)?; - ebml::end_master(&mut self.writer, adds_pos)?; - ebml::end_master(&mut self.writer, bg_pos)?; - Ok(()) + let mut buf = std::mem::take(&mut self.block_group_buf); + buf.clear(); + let res = Self::build_block_group( + &mut buf, + track_num, + relative_ts, + data, + reference, + duration_ticks, + Some(additional), + ) + .and_then(|()| self.writer.write_all(&buf)); + self.block_group_buf = buf; + res } }