direct io for blocks

10 files changed, 61 insertions, 53 deletions

diff --git a/ebml_derive/src/lib.rs b/ebml_derive/src/lib.rs
index ab6cff8..fdf04e8 100644
--- a/ebml_derive/src/lib.rs
+++ b/ebml_derive/src/lib.rs
@@ -35,6 +35,7 @@ pub fn define_ebml(ts: TokenStream) -> TokenStream {
                         "Float" => quote!((f64)),
                         "Utf8" => quote!((String)),
                         "Binary" => quote!((Vec<u8>)),
+                        "Block" => quote!((Block)),
                         _ => panic!("unsupported type {type}"),
                     },
                 })
@@ -110,6 +111,7 @@ pub fn define_ebml(ts: TokenStream) -> TokenStream {
     quote! {
         use crate::Master;
         use crate::WriteValue;
+        use crate::Block;
 
         #[derive(Debug, PartialEq, Clone)]
         pub enum MatroskaTag {
diff --git a/matroska/src/bin/mkvdump.rs b/matroska/src/bin/mkvdump.rs
index 27b4849..b58adcc 100644
--- a/matroska/src/bin/mkvdump.rs
+++ b/matroska/src/bin/mkvdump.rs
@@ -3,7 +3,7 @@
     which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
     Copyright (C) 2024 metamuffin <metamuffin.org>
 */
-use jellymatroska::{block::Block, matroska::MatroskaTag, read::EbmlReader};
+use jellymatroska::{matroska::MatroskaTag, read::EbmlReader};
 use std::{fs::File, io::BufReader};
 
 fn main() {
@@ -15,7 +15,6 @@ fn main() {
         let (position, tag) = tag.unwrap();
         match tag {
             MatroskaTag::SimpleBlock(b) | MatroskaTag::Block(b) => {
-                let b = Block::parse(&b).unwrap();
                 println!("block kf={} ts_off={}", b.keyframe, b.timestamp_off)
             }
             _ => println!("{} {tag:?}", position.unwrap_or(0)),
diff --git a/matroska/src/block.rs b/matroska/src/block.rs
index 5ab398b..1ab8ceb 100644
--- a/matroska/src/block.rs
+++ b/matroska/src/block.rs
@@ -3,16 +3,19 @@
     which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
     Copyright (C) 2024 metamuffin <metamuffin.org>
 */
+use crate::write::vint_length;
 use crate::{read::ReadExt, write::write_vint};
-use crate::Result;
-use std::io::Cursor;
+use crate::{ReadValue, Result, WriteValue};
+use std::io::{Cursor, Write};
 
+#[derive(Debug, PartialEq, Clone, Copy)]
 pub enum LacingType {
     Xiph,
     FixedSize,
     Ebml,
 }
 
+#[derive(Debug, PartialEq, Clone)]
 pub struct Block {
     pub track: u64,
     pub keyframe: bool,
@@ -23,8 +26,8 @@ pub struct Block {
     pub data: Vec<u8>,
 }
 
-impl Block {
-    pub fn parse(buf: &[u8]) -> Result<Self> {
+impl ReadValue for Block {
+    fn from_buf(buf: &[u8]) -> Result<Self> {
         let (track, c) = Cursor::new(buf).read_vint_len()?;
         let timestamp_off = i16::from_be_bytes(buf[c..c + 2].try_into().unwrap());
         let flags = buf[c + 2];
@@ -51,28 +54,42 @@ impl Block {
             timestamp_off,
         })
     }
-    pub fn dump(&self) -> Vec<u8> {
-        let mut out = vec![];
-        write_vint(&mut out, self.track).unwrap();
-        out.extend(self.timestamp_off.to_be_bytes().into_iter());
-        out.push(
-            match self.discardable {
-                true => 0b1,
-                false => 0b0,
-            } | match self.keyframe {
-                true => 0b10000000,
-                false => 0b00000000,
-            } | match self.invisible {
-                true => 0b1000,
-                false => 0b0000,
-            } | match self.lacing {
-                Some(LacingType::Xiph) => 0b010,
-                Some(LacingType::Ebml) => 0b100,
-                Some(LacingType::FixedSize) => 0b110,
-                None => 0b0000,
-            },
-        );
-        out.extend(self.data.iter());
-        out
+}
+
+impl WriteValue for Block {
+    fn write_to(&self, w: &mut impl Write) -> Result<()> {
+        write_vint(w, self.inner_len() as u64)?;
+        write_vint(w, self.track)?;
+        w.write_all(&self.timestamp_off.to_be_bytes())?;
+        w.write_all(&[match self.discardable {
+            true => 0b1,
+            false => 0b0,
+        } | match self.keyframe {
+            true => 0b10000000,
+            false => 0b00000000,
+        } | match self.invisible {
+            true => 0b1000,
+            false => 0b0000,
+        } | match self.lacing {
+            Some(LacingType::Xiph) => 0b010,
+            Some(LacingType::Ebml) => 0b100,
+            Some(LacingType::FixedSize) => 0b110,
+            None => 0b0000,
+        }])?;
+        w.write_all(&self.data)?;
+        Ok(())
+    }
+    fn size(&self) -> usize {
+        let il = self.inner_len();
+        vint_length(il as u64) + il
+    }
+}
+
+impl Block {
+    fn inner_len(&self) -> usize {
+        vint_length(self.track)
+        + 2 // timestamp
+        + 1 // flags
+        + self.data.len()
     }
 }
diff --git a/matroska/src/lib.rs b/matroska/src/lib.rs
index 09635d4..7e4121b 100644
--- a/matroska/src/lib.rs
+++ b/matroska/src/lib.rs
@@ -22,5 +22,6 @@ pub enum Master {
     End,
 }
 
+pub(crate) use block::Block;
 pub(crate) use error::Error;
 pub(crate) type Result<T> = core::result::Result<T, Error>;
diff --git a/matroska/src/matroska.rs b/matroska/src/matroska.rs
index 6bbc5f7..937fca7 100644
--- a/matroska/src/matroska.rs
+++ b/matroska/src/matroska.rs
@@ -76,7 +76,7 @@ define_ebml! {
 
         Cluster[0x1F43B675]: {
             BlockGroup[0xA0]: {
-                Block[0xA1]: Binary,
+                Block[0xA1]: Block,
                 BlockAdditions[0x75A1]: {
                     BlockMore[0xA6]: {
                         BlockAddID[0xEE]: Uint,
@@ -110,7 +110,7 @@ define_ebml! {
             SilentTracks[0x5854]: {
                 SilentTrackNumber[0x58D7]: Uint,
             },
-            SimpleBlock[0xA3]: Binary,
+            SimpleBlock[0xA3]: Block,
             Timestamp[0xE7]: Uint,
         },
 
diff --git a/remuxer/src/extract.rs b/remuxer/src/extract.rs
index 1b2b50c..bef5d4d 100644
--- a/remuxer/src/extract.rs
+++ b/remuxer/src/extract.rs
@@ -44,14 +44,13 @@ pub fn read_group(segment: &mut EbmlReader) -> anyhow::Result<(u64, Block)> {
             MatroskaTag::Crc32(_) => (),
             MatroskaTag::Cluster(_) => bail!("unexpected cluster"),
             MatroskaTag::Timestamp(_) => (),
-            MatroskaTag::SimpleBlock(_buf) => {
-                let block = Block::parse(&_buf)?;
+            MatroskaTag::SimpleBlock(block) => {
                 return Ok((1000, block)); // HDMV/PGS does not use duration?!
             }
             MatroskaTag::BlockGroup(Master::Start) => (),
             MatroskaTag::BlockGroup(Master::End) => return Ok((dur.unwrap(), block.unwrap())),
             MatroskaTag::BlockDuration(duration) => dur = Some(duration),
-            MatroskaTag::Block(buf) => block = Some(Block::parse(&buf)?),
+            MatroskaTag::Block(blk) => block = Some(blk),
             MatroskaTag::Cues(_) => bail!("reached cues, this is the end"),
             MatroskaTag::Segment(Master::End) => bail!("extractor reached segment end"),
             _ => debug!("(rs) tag ignored: {item:?}"),
diff --git a/remuxer/src/remux.rs b/remuxer/src/remux.rs
index 851b43d..a919354 100644
--- a/remuxer/src/remux.rs
+++ b/remuxer/src/remux.rs
@@ -13,7 +13,6 @@ use jellycommon::{
     LocalTrack, NodePublic, SourceTrack,
 };
 use jellymatroska::{
-    block::Block,
     read::EbmlReader,
     write::{bad_vint_length, vint_length, EbmlWriter},
     Master, MatroskaTag,
@@ -262,7 +261,6 @@ pub fn remux_stream_into(
     }
 
     struct ReaderD<'a> {
-        peek: Option<Block>,
         stream: SegmentExtractIter<'a>,
         mapped: u64,
     }
@@ -280,12 +278,10 @@ pub fn remux_stream_into(
                     MatroskaTag::Cluster(Master::Start), // TODO shouldn't this be a child of cluster?
                 )
                 .context("seeking in input")?;
-            let mut stream =
-                SegmentExtractIter::new(&mut inp.reader, inp.source_track_index as u64);
+            let stream = SegmentExtractIter::new(&mut inp.reader, inp.source_track_index as u64);
 
             Ok(ReaderD {
                 mapped: inp.mapped,
-                peek: Some(stream.next()?.0), // TODO handle duration
                 stream,
             })
         })
@@ -309,10 +305,8 @@ pub fn remux_stream_into(
         let mut cluster_blocks = vec![MatroskaTag::Timestamp(cluster.timestamp)];
         for (block_track, index_block) in cluster.blocks {
             let track_reader = &mut track_readers[block_track];
-            let mut block = track_reader
-                .peek
-                .replace(track_reader.stream.next()?.0) // TODO handle duration
-                .expect("source file too short");
+            // TODO handle duration
+            let mut block = track_reader.stream.next()?.0;
 
             assert_eq!(index_block.size, block.data.len(), "seek index is wrong");
 
@@ -320,8 +314,7 @@ pub fn remux_stream_into(
             block.timestamp_off = (index_block.pts - cluster.timestamp).try_into().unwrap();
             trace!("n={} tso={}", block.track, block.timestamp_off);
 
-            let buf = block.dump();
-            cluster_blocks.push(MatroskaTag::SimpleBlock(buf))
+            cluster_blocks.push(MatroskaTag::SimpleBlock(block))
         }
         output.write_tag(&MatroskaTag::Cluster(Master::Collected(cluster_blocks)))?;
     }
diff --git a/remuxer/src/seek_index.rs b/remuxer/src/seek_index.rs
index 7e9cee1..5bec6e6 100644
--- a/remuxer/src/seek_index.rs
+++ b/remuxer/src/seek_index.rs
@@ -84,8 +84,7 @@ fn import_seek_index_segment(
                                 })) = children.n()
                                 {
                                     match item {
-                                        MatroskaTag::Block(ref buf) => {
-                                            let block = Block::parse(buf)?;
+                                        MatroskaTag::Block(ref block) => {
                                             debug!(
                                                 "block: track={} tso={}",
                                                 block.track, block.timestamp_off
@@ -101,8 +100,7 @@ fn import_seek_index_segment(
                                     }
                                 }
                             }
-                            MatroskaTag::SimpleBlock(buf) => {
-                                let block = Block::parse(&buf)?;
+                            MatroskaTag::SimpleBlock(block) => {
                                 trace!(
                                     "simple block: track={} tso={}",
                                     block.track,
diff --git a/remuxer/src/segment_extractor.rs b/remuxer/src/segment_extractor.rs
index ec645c3..2a12802 100644
--- a/remuxer/src/segment_extractor.rs
+++ b/remuxer/src/segment_extractor.rs
@@ -40,8 +40,7 @@ impl<'a> SegmentExtractIter<'a> {
                     group = false;
                 }
                 MatroskaTag::BlockDuration(d) => duration = Some(d),
-                MatroskaTag::SimpleBlock(buf) | MatroskaTag::Block(buf) => {
-                    let block = Block::parse(&buf)?;
+                MatroskaTag::SimpleBlock(block) | MatroskaTag::Block(block) => {
                     if block.track == self.extract {
                         trace!("block: track={} tso={}", block.track, block.timestamp_off);
                         if group {
diff --git a/remuxer/src/snippet.rs b/remuxer/src/snippet.rs
index daa50a9..ff391a2 100644
--- a/remuxer/src/snippet.rs
+++ b/remuxer/src/snippet.rs
@@ -196,10 +196,10 @@ pub fn write_snippet_into(
             if let Some(duration) = duration {
                 blocks.push(MatroskaTag::BlockGroup(Master::Collected(vec![
                     MatroskaTag::BlockDuration(duration),
-                    MatroskaTag::Block(block.dump()),
+                    MatroskaTag::Block(block),
                 ])))
             } else {
-                blocks.push(MatroskaTag::SimpleBlock(block.dump()))
+                blocks.push(MatroskaTag::SimpleBlock(block))
             }
         }
         output.write_tag(&MatroskaTag::Cluster(Master::Collected(blocks)))?;