less useless copying in the matroska library

4 files changed, 177 insertions, 22 deletions

diff --git a/ebml_derive/src/lib.rs b/ebml_derive/src/lib.rs
index cd2d194..34510b9 100644
--- a/ebml_derive/src/lib.rs
+++ b/ebml_derive/src/lib.rs
@@ -81,6 +81,10 @@ pub fn define_ebml(ts: TokenStream) -> TokenStream {
         .iter()
         .map(|Tag { name, .. }| quote! { Self::#name(v) => v.write_to(w) })
         .collect::<Vec<_>>();
+    let size_match = tags
+        .iter()
+        .map(|Tag { name, .. }| quote! { Self::#name(v) => v.size() })
+        .collect::<Vec<_>>();
     let cons_master_match = tags
         .iter()
         .filter_map(
@@ -131,6 +135,9 @@ pub fn define_ebml(ts: TokenStream) -> TokenStream {
             pub fn write(&self, w: &mut Vec<u8>) -> crate::Result<()> {
                 match self { #(#write_match),* }
             }
+            pub fn size(&self) -> usize {
+                match self { #(#size_match),* }
+            }
         }
     }
     .into()
diff --git a/matroska/src/write.rs b/matroska/src/write.rs
index fe6033f..3b127f8 100644
--- a/matroska/src/write.rs
+++ b/matroska/src/write.rs
@@ -91,13 +91,25 @@ impl<W: Seek> Seek for EbmlWriter<W> {
 
 impl MatroskaTag {
     pub fn write_full(&self, w: &mut Vec<u8>) -> Result<()> {
-        let mut buf = vec![];
-        buf.extend(self.id().to_be_bytes().iter().skip_while(|&v| *v == 0u8));
-        //* note: it is relevant here, to pass the buffer with the id, such that closing tags, can clear it
-        self.write(&mut buf)?;
-        w.extend_from_slice(&buf);
+        w.extend(self.id().to_be_bytes().iter().skip_while(|&v| *v == 0u8));
+        self.write(w)?;
         Ok(())
     }
+    pub fn size_full(&self) -> usize {
+        (8 - self.id().leading_zeros() as usize / 8) + self.size()
+    }
+}
+
+pub fn write_vint(w: &mut Vec<u8>, i: u64) -> Result<()> {
+    if i > (1 << 56) - 1 {
+        Err(Error::VarintTooLong)?
+    }
+    let len = (64 - i.leading_zeros() as usize) / 7 + 1;
+    let mut bytes = i.to_be_bytes();
+    let trunc = &mut bytes[(8 - len)..];
+    trunc[0] |= 1 << (8 - len);
+    w.extend_from_slice(trunc);
+    Ok(())
 }
 
 /// this routine works only, if the varint is as small as it can possibly be.
@@ -124,6 +136,7 @@ pub fn bad_vint_length(v: u64) -> usize {
 pub trait WriteValue {
     /// writes the contents of a tag, including the size but excluding the id.
     fn write_to(&self, w: &mut Vec<u8>) -> Result<()>;
+    fn size(&self) -> usize;
 }
 
 impl WriteValue for i64 {
@@ -148,6 +161,15 @@ impl WriteValue for i64 {
         };
         Ok(())
     }
+
+    fn size(&self) -> usize {
+        1 + match 64 - self.leading_zeros() {
+            x if x <= 8 => 1,
+            x if x <= 16 => 2,
+            x if x <= 32 => 4,
+            _ => 8,
+        }
+    }
 }
 impl WriteValue for u64 {
     fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
@@ -171,6 +193,14 @@ impl WriteValue for u64 {
         };
         Ok(())
     }
+    fn size(&self) -> usize {
+        1 + match 64 - self.leading_zeros() {
+            x if x <= 8 => 1,
+            x if x <= 16 => 2,
+            x if x <= 32 => 4,
+            _ => 8,
+        }
+    }
 }
 impl WriteValue for f64 {
     fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
@@ -178,6 +208,9 @@ impl WriteValue for f64 {
         w.extend_from_slice(&self.to_be_bytes());
         Ok(())
     }
+    fn size(&self) -> usize {
+        1 + 8
+    }
 }
 impl WriteValue for Vec<u8> {
     fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
@@ -185,6 +218,10 @@ impl WriteValue for Vec<u8> {
         w.extend_from_slice(self);
         Ok(())
     }
+
+    fn size(&self) -> usize {
+        vint_length(self.len() as u64) + self.len()
+    }
 }
 impl WriteValue for String {
     fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
@@ -193,6 +230,10 @@ impl WriteValue for String {
         w.extend_from_slice(sl);
         Ok(())
     }
+
+    fn size(&self) -> usize {
+        vint_length(self.as_bytes().len() as u64) + self.as_bytes().len()
+    }
 }
 impl WriteValue for EbmlSize {
     fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
@@ -202,6 +243,13 @@ impl WriteValue for EbmlSize {
         }
         Ok(())
     }
+
+    fn size(&self) -> usize {
+        match self {
+            EbmlSize::Exact(s) => vint_length(*s as u64),
+            EbmlSize::Unknown => 8,
+        }
+    }
 }
 
 impl WriteValue for Master {
@@ -209,30 +257,130 @@ impl WriteValue for Master {
         match self {
             Master::Start => EbmlSize::Unknown.write_to(w),
             Master::End => {
-                w.clear();
-                Ok(())
+                unreachable!()
             }
             Master::Collected(c) => {
-                let mut ib = vec![];
+                let mut size = 0;
+                for c in c {
+                    size += c.size_full();
+                }
+                EbmlSize::Exact(size).write_to(w)?;
+                let k = w.len();
                 for c in c {
-                    c.write_full(&mut ib)?;
+                    c.write_full(w)?;
                 }
-                EbmlSize::Exact(ib.len()).write_to(w)?;
-                w.extend_from_slice(&ib);
+                debug_assert_eq!(w.len() - k, size, "real size of collected differs");
                 Ok(())
             }
         }
     }
+    fn size(&self) -> usize {
+        match self {
+            Master::Start => EbmlSize::Unknown.size(),
+            Master::End => unreachable!(),
+            Master::Collected(c) => {
+                let mut size = 0;
+                for c in c {
+                    size += c.size_full();
+                }
+                EbmlSize::Exact(size).size() + size
+            }
+        }
+    }
 }
 
-pub fn write_vint(w: &mut Vec<u8>, i: u64) -> Result<()> {
-    if i > (1 << 56) - 1 {
-        Err(Error::VarintTooLong)?
+#[cfg(test)]
+mod test {
+    use crate::{Master, MatroskaTag, WriteValue};
+
+    #[test]
+    fn test_int_size() {
+        let test = |x: i64| {
+            eprintln!("{x:?}");
+            let mut out = Vec::new();
+            x.write_to(&mut out).unwrap();
+            assert_eq!(out.len(), x.size())
+        };
+        test(1);
+        test(2);
+        test(20);
+        test(200);
+        test(2000);
+        test(20000);
+        test(200000);
+    }
+
+    #[test]
+    fn test_uint_size() {
+        let test = |x: u64| {
+            eprintln!("{x:?}");
+            let mut out = Vec::new();
+            x.write_to(&mut out).unwrap();
+            assert_eq!(out.len(), x.size())
+        };
+        test(1);
+        test(2);
+        test(20);
+        test(200);
+        test(2000);
+        test(20000);
+        test(200000);
+    }
+
+    #[test]
+    fn test_string_size() {
+        let test = |x: &str| {
+            eprintln!("{x:?}");
+            let x = x.to_owned();
+            let mut out = Vec::new();
+            x.write_to(&mut out).unwrap();
+            assert_eq!(out.len(), x.size())
+        };
+        test("");
+        test("x");
+        test("wub");
+        test("Hello world");
+        test("just making sure that");
+        test("this is actually working *exactly* how i want it to");
+        test("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa");
+    }
+
+    #[test]
+    fn test_vec_size() {
+        let test = |x: &[u8]| {
+            eprintln!("{x:?}");
+            let x = x.to_owned();
+            let mut out = Vec::new();
+            x.write_to(&mut out).unwrap();
+            assert_eq!(out.len(), x.size())
+        };
+        test(&[]);
+        test(&[1]);
+        test(&[1, 2]);
+        test(&[23, 4, 4, 23, 4, 234, 232, 4, 234, 23, 1]);
+        test(&[
+            34, 123, 5, 1, 3, 13, 1, 23, 12, 5, 5, 3, 123, 12, 3, 13, 12, 5, 3, 123, 13, 1, 3,
+        ]);
+    }
+
+    #[test]
+    fn test_master_size() {
+        let test = |x: Master| {
+            eprintln!("{x:?}");
+            let x = x.to_owned();
+            let mut out = Vec::new();
+            x.write_to(&mut out).unwrap();
+            assert_eq!(out.len(), x.size())
+        };
+        test(Master::Start);
+        // test(Master::End);
+        test(Master::Collected(vec![]));
+        test(Master::Collected(vec![MatroskaTag::EbmlVersion(1)]));
+        test(Master::Collected(vec![
+            MatroskaTag::EbmlVersion(1),
+            MatroskaTag::EbmlMaxSizeLength(4),
+            MatroskaTag::EbmlReadVersion(3),
+            MatroskaTag::EbmlMaxIdLength(4),
+        ]));
     }
-    let len = (64 - i.leading_zeros() as usize) / 7 + 1;
-    let mut bytes = i.to_be_bytes();
-    let trunc = &mut bytes[(8 - len)..];
-    trunc[0] |= 1 << (8 - len);
-    w.extend_from_slice(trunc);
-    Ok(())
 }
diff --git a/remuxer/src/remux.rs b/remuxer/src/remux.rs
index 1e75b06..8273c9e 100644
--- a/remuxer/src/remux.rs
+++ b/remuxer/src/remux.rs
@@ -325,7 +325,7 @@ pub fn remux_stream_into(
         }
         output.write_tag(&MatroskaTag::Cluster(Master::Collected(cluster_blocks)))?;
     }
-    output.write_tag(&MatroskaTag::Segment(Master::End))?;
+    // output.write_tag(&MatroskaTag::Segment(Master::End))?;
     Ok(())
 }
 
diff --git a/remuxer/src/snippet.rs b/remuxer/src/snippet.rs
index 8fcb109..75e0471 100644
--- a/remuxer/src/snippet.rs
+++ b/remuxer/src/snippet.rs
@@ -198,7 +198,7 @@ pub fn write_snippet_into(
         output.write_tag(&MatroskaTag::Cluster(Master::Collected(blocks)))?;
     }
 
-    output.write_tag(&MatroskaTag::Segment(Master::End))?;
+    // output.write_tag(&MatroskaTag::Segment(Master::End))?;
 
     debug!("wrote {} bytes", output.position());
     Ok(())