write more code

1 files changed, 169 insertions, 0 deletions

diff --git a/matroska/src/write.rs b/matroska/src/write.rs
new file mode 100644
index 0000000..fc12ffc
--- /dev/null
+++ b/matroska/src/write.rs
@@ -0,0 +1,169 @@
+use anyhow::{bail, Result};
+use std::io::Write;
+
+use crate::{matroska::MatroskaTag, size::EbmlSize, Master};
+
+pub struct EbmlWriter {
+    inner: Box<dyn Write>,
+    position: usize,
+}
+
+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<()> {
+        self.inner.write_all(data)?;
+        self.position += data.len();
+        Ok(())
+    }
+
+    pub fn write_tag(&mut self, tag: &MatroskaTag) -> Result<()> {
+        let mut buf = vec![];
+        tag.write_full(&mut buf)?;
+        self.write(&buf)?;
+        Ok(())
+    }
+
+    pub fn write_vint(&mut self, 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);
+        self.write(&trunc)
+    }
+}
+
+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);
+        Ok(())
+    }
+}
+
+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 => EbmlSize::Unknown.write_to(w),
+            Master::End => Ok(w.clear()),
+            Master::Collected(c) => {
+                let mut ib = vec![];
+                for c in c {
+                    c.write_full(&mut ib)?;
+                }
+                EbmlSize::Exact(ib.len()).write_to(w)?;
+                w.extend_from_slice(&ib);
+                Ok(())
+            }
+        }
+    }
+}
+
+pub fn write_vint(w: &mut Vec<u8>, i: u64) -> Result<()> {
+    if i > (1 << 56) - 1 {
+        bail!("vint does not fit");
+    }
+    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(())
+}