seeking and broken writing

1 files changed, 134 insertions, 7 deletions

diff --git a/ebml/src/write.rs b/ebml/src/write.rs
index b294d44..f57a5c9 100644
--- a/ebml/src/write.rs
+++ b/ebml/src/write.rs
@@ -1,14 +1,19 @@
 use anyhow::{bail, Result};
 use std::io::Write;
 
-pub struct EbmlWriter<T> {
-    inner: T,
+use crate::{matroska::MatroskaTag, size::EbmlSize, Master};
+
+pub struct EbmlWriter {
+    inner: Box<dyn Write>,
     position: usize,
 }
 
-impl<T: Write> EbmlWriter<T> {
-    pub fn new(inner: T, position: usize) -> Self {
-        Self { inner, position }
+impl EbmlWriter {
+    pub fn new<T: Write + 'static>(inner: T, position: usize) -> Self {
+        Self {
+            inner: Box::new(inner),
+            position,
+        }
     }
 
     pub fn write(&mut self, data: &[u8]) -> Result<()> {
@@ -17,6 +22,21 @@ impl<T: Write> EbmlWriter<T> {
         Ok(())
     }
 
+    pub fn write_tag(&mut self, tag: &MatroskaTag) -> Result<()> {
+        self.write_tag_id(tag.id())?;
+        let mut buf = vec![];
+        tag.write(&mut buf)?;
+        self.write(&buf)?;
+        Ok(())
+    }
+
+    pub fn write_tag_id(&mut self, id: u64) -> Result<()> {
+        for n in id.to_be_bytes().iter().skip_while(|&v| *v == 0u8) {
+            self.write(&[*n])?;
+        }
+        Ok(())
+    }
+
     pub fn write_vint(&mut self, i: u64) -> Result<()> {
         if i > (1 << 56) - 1 {
             bail!("vint does not fit");
@@ -29,8 +49,115 @@ impl<T: Write> EbmlWriter<T> {
             len += 1;
         }
         let mut bytes = i.to_be_bytes();
-        bytes[0] |= 1 << (8 - len);
-        self.write(&bytes)
+        let trunc = &mut bytes[(8 - len)..];
+        trunc[0] |= 1 << (8 - len);
+        self.write(&trunc)
     }
 }
 
+pub trait WriteValue {
+    fn write_to(&self, w: &mut Vec<u8>) -> Result<()>;
+}
+
+impl WriteValue for i64 {
+    fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
+        Ok(match 64 - self.leading_zeros() {
+            x if x <= 8 => {
+                w.push(0x81);
+                w.extend_from_slice(&(*self as i8).to_be_bytes());
+            }
+            x if x <= 16 => {
+                w.push(0x82);
+                w.extend_from_slice(&(*self as i16).to_be_bytes());
+            }
+            x if x <= 32 => {
+                w.push(0x84);
+                w.extend_from_slice(&(*self as i32).to_be_bytes());
+            }
+            _ => {
+                w.push(0x88);
+                w.extend_from_slice(&self.to_be_bytes());
+            }
+        })
+    }
+}
+impl WriteValue for u64 {
+    fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
+        Ok(match 64 - self.leading_zeros() {
+            x if x <= 8 => {
+                w.push(0x81);
+                w.extend_from_slice(&(*self as u8).to_be_bytes());
+            }
+            x if x <= 16 => {
+                w.push(0x82);
+                w.extend_from_slice(&(*self as u16).to_be_bytes());
+            }
+            x if x <= 32 => {
+                w.push(0x84);
+                w.extend_from_slice(&(*self as u32).to_be_bytes());
+            }
+            _ => {
+                w.push(0x88);
+                w.extend_from_slice(&self.to_be_bytes());
+            }
+        })
+    }
+}
+impl WriteValue for f64 {
+    fn write_to(&self, w: &mut Vec<u8>) -> Result<(), anyhow::Error> {
+        w.push(0x88);
+        w.extend_from_slice(&self.to_be_bytes());
+        Ok(())
+    }
+}
+impl WriteValue for Vec<u8> {
+    fn write_to(&self, w: &mut Vec<u8>) -> Result<(), anyhow::Error> {
+        write_vint(w, self.len() as u64)?;
+        w.extend_from_slice(&self);
+        Ok(())
+    }
+}
+impl WriteValue for String {
+    fn write_to(&self, w: &mut Vec<u8>) -> Result<(), anyhow::Error> {
+        let sl = self.as_bytes();
+        write_vint(w, sl.len() as u64)?;
+        w.extend_from_slice(sl);
+        Ok(())
+    }
+}
+impl WriteValue for EbmlSize {
+    fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
+        match self {
+            EbmlSize::Exact(s) => write_vint(w, *s as u64)?,
+            EbmlSize::Unknown => w.extend_from_slice(&(u64::MAX >> 7).to_be_bytes()),
+        }
+        Ok(())
+    }
+}
+
+impl WriteValue for Master {
+    fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
+        match self {
+            Master::Start(size) => size.write_to(w),
+            Master::End => Ok(()),
+        }
+    }
+}
+
+pub fn write_vint(w: &mut Vec<u8>, i: u64) -> Result<()> {
+    if i > (1 << 56) - 1 {
+        bail!("vint does not fit");
+    }
+    let mut len = 1;
+    while len <= 8 {
+        if i < (1 << ((7 * len) - 1)) {
+            break;
+        }
+        len += 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(())
+}