path: root/matroska/src/write.rs

/*
    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
    Copyright (C) 2024 metamuffin <metamuffin.org>
*/
use crate::error::Error;
use crate::Result;
use crate::{matroska::MatroskaTag, size::EbmlSize, Master};
use log::debug;
use std::io::{Seek, Write};

pub struct EbmlWriter<W> {
    inner: W,
    position: usize,
}

impl<W: Write> EbmlWriter<W> {
    pub fn new(inner: W, position: usize) -> Self {
        Self { inner, position }
    }

    pub fn write(&mut self, data: &[u8]) -> Result<()> {
        self.inner.write_all(data).map_err(Error::Io)?;
        self.position += data.len();
        Ok(())
    }

    pub fn position(&self) -> usize {
        self.position
    }

    pub fn write_padding(&mut self, position: usize) -> Result<()> {
        debug!("padding up to {position}");
        let mut size = position - self.position;
        match size {
            0 => return Ok(()),
            1 => Err(Error::InvalidPadding)?,
            _ => (),
        }
        size -= 1; // subtract tag size
        size -= 4; // subtract vint size

        // match size {
        //     _ if size < (1 << 7) => size -= 1,
        //     _ if size < (1 << 14) => size -= 2,
        //     _ if size < (1 << 21) => size -= 3,
        //     _ if size < (1 << 28) => size -= 4,
        //     _ if size < (1 << 35) => size -= 5,
        //     _ => bail!("padding to large"),
        // }

        self.write(&[0xec])?;
        self.write_vint_len(size.try_into().unwrap(), 4)?;
        self.write(&vec![0; size])?;
        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 {
            Err(Error::VarintTooLong)?
        }
        self.write_vint_len(i, vint_length(i))
    }

    pub fn write_vint_len(&mut self, i: u64, len: usize) -> Result<()> {
        let mut bytes = i.to_be_bytes();
        let trunc = &mut bytes[(8 - len)..];
        trunc[0] |= 1 << (8 - len);
        self.write(trunc)
    }
}

impl<W: Seek> Seek for EbmlWriter<W> {
    fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> {
        self.inner.seek(pos)?;
        match pos {
            std::io::SeekFrom::Start(s) => self.position = s as usize,
            std::io::SeekFrom::End(_) => unimplemented!(),
            std::io::SeekFrom::Current(s) => self.position += s as usize,
        }
        Ok(self.position as u64)
    }
}

impl MatroskaTag {
    pub fn write_full(&self, w: &mut Vec<u8>) -> Result<()> {
        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.
/// thats not always what we do though - see below
pub fn vint_length(v: u64) -> usize {
    let mut len = 1;
    while len <= 8 {
        if v < (1 << ((7 * len) - 1)) {
            break;
        }
        len += 1;
    }
    len
}
pub fn bad_vint_length(v: u64) -> usize {
    match 64 - v.leading_zeros() {
        x if x <= 8 => 1,
        x if x <= 16 => 2,
        x if x <= 32 => 4,
        _ => 8,
    }
}

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 {
    fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
        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());
            }
        };
        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<()> {
        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());
            }
        };
        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<()> {
        w.push(0x88);
        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<()> {
        write_vint(w, self.len() as u64)?;
        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<()> {
        let sl = self.as_bytes();
        write_vint(w, sl.len() as u64)?;
        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<()> {
        match self {
            EbmlSize::Exact(s) => write_vint(w, *s as u64)?,
            EbmlSize::Unknown => w.extend_from_slice(&(u64::MAX >> 7).to_be_bytes()),
        }
        Ok(())
    }

    fn size(&self) -> usize {
        match self {
            EbmlSize::Exact(s) => vint_length(*s as u64),
            EbmlSize::Unknown => 8,
        }
    }
}

impl WriteValue for Master {
    fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
        match self {
            Master::Start => EbmlSize::Unknown.write_to(w),
            Master::End => {
                unreachable!()
            }
            Master::Collected(c) => {
                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(w)?;
                }
                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
            }
        }
    }
}

#[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),
        ]));
    }
}