avif_serialize/
writer.rs

1use std::convert::TryFrom;
2use std::io;
3
4pub struct OOM;
5
6pub trait WriterBackend {
7    type Error;
8    fn reserve(&mut self, size: usize) -> Result<(), Self::Error>;
9    fn extend_from_slice_in_capacity(&mut self, data: &[u8]) -> Result<(), Self::Error>;
10}
11
12/// `io::Write` generates bloated code (with backtrace for every byte written),
13/// so small boxes are written infallibly.
14impl WriterBackend for Vec<u8> {
15    type Error = OOM;
16
17    #[inline]
18    fn reserve(&mut self, size: usize) -> Result<(), Self::Error> {
19        self.try_reserve(size).map_err(|_| OOM)
20    }
21
22    #[inline(always)]
23    fn extend_from_slice_in_capacity(&mut self, data: &[u8]) -> Result<(), Self::Error> {
24        let has_capacity = self.capacity() - self.len() >= data.len();
25        debug_assert!(has_capacity);
26        if has_capacity {
27            self.extend_from_slice(data);
28            Ok(())
29        } else {
30            Err(OOM)
31        }
32    }
33}
34
35pub struct IO<W>(pub W);
36
37impl<W: io::Write> WriterBackend for IO<W> {
38    type Error = io::Error;
39
40    #[inline]
41    fn reserve(&mut self, _size: usize) -> io::Result<()> {
42        Ok(())
43    }
44
45    #[inline(always)]
46    fn extend_from_slice_in_capacity(&mut self, data: &[u8]) -> io::Result<()> {
47        self.0.write_all(data)
48    }
49}
50
51pub struct Writer<'p, 'w, B> {
52    out: &'w mut B,
53    #[cfg(debug_assertions)]
54    parent: Option<&'p mut usize>,
55    #[cfg(not(debug_assertions))]
56    parent: std::marker::PhantomData<&'p mut usize>,
57    #[cfg(debug_assertions)]
58    left: usize,
59}
60
61impl<'w, B> Writer<'static, 'w, B> {
62    #[inline]
63    pub fn new(out: &'w mut B) -> Self {
64        Self {
65            parent: Default::default(),
66            #[cfg(debug_assertions)]
67            left: 0,
68            out,
69        }
70    }
71}
72
73impl<B: WriterBackend> Writer<'_, '_, B> {
74    #[inline(always)]
75    pub fn full_box(&mut self, len: usize, typ: [u8; 4], version: u8) -> Result<Writer<'_, '_, B>, B::Error> {
76        let mut b = self.basic_box(len, typ)?;
77        b.push(&[version, 0, 0, 0])?;
78        Ok(b)
79    }
80
81    #[inline]
82    pub fn basic_box(&mut self, len: usize, typ: [u8; 4]) -> Result<Writer<'_, '_, B>, B::Error> {
83        let mut b = Writer {
84            out: self.out,
85            parent: Default::default(),
86            #[cfg(debug_assertions)]
87            left: len,
88        };
89        #[cfg(debug_assertions)]
90        if self.left > 0 {
91            self.left -= len;
92            b.parent = Some(&mut self.left);
93        } else {
94            debug_assert!(self.parent.is_none());
95        }
96        b.out.reserve(len)?;
97
98        if let Ok(len) = u32::try_from(len) {
99            b.u32(len)?;
100        } else {
101            debug_assert!(false, "constants for box size don't include this");
102            b.u32(1)?;
103            b.u64(len as u64)?;
104        }
105        b.push(&typ)?;
106        Ok(b)
107    }
108
109    #[inline(always)]
110    pub fn push(&mut self, data: &[u8]) -> Result<(), B::Error> {
111        #[cfg(debug_assertions)] {
112            self.left -= data.len();
113        }
114        self.out.extend_from_slice_in_capacity(data)
115    }
116
117    #[inline(always)]
118    pub fn u8(&mut self, val: u8) -> Result<(), B::Error> {
119        self.push(std::slice::from_ref(&val))
120    }
121
122    #[inline(always)]
123    pub fn u16(&mut self, val: u16) -> Result<(), B::Error> {
124        self.push(&val.to_be_bytes())
125    }
126
127    #[inline(always)]
128    pub fn u32(&mut self, val: u32) -> Result<(), B::Error> {
129        self.push(&val.to_be_bytes())
130    }
131
132    #[inline(always)]
133    pub fn u64(&mut self, val: u64) -> Result<(), B::Error> {
134        self.push(&val.to_be_bytes())
135    }
136}
137
138#[cfg(debug_assertions)]
139impl<B> Drop for Writer<'_, '_, B> {
140    fn drop(&mut self) {
141        assert_eq!(self.left, 0);
142    }
143}