Skip to main content

mkv_element/
base.rs

1use crate::error::Error;
2use crate::io::blocking_impl::*;
3
4use crate::*;
5
6use std::fmt::Debug;
7use std::fmt::Display;
8use std::ops::Deref;
9
10/// A variable-length integer RFC 8794
11#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
12pub struct VInt64 {
13    /// The decoded integer value.
14    pub value: u64,
15    /// Whether this VInt64 represents an unknown size.
16    pub is_unknown: bool,
17}
18
19impl Display for VInt64 {
20    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        // write!(f, "{}", self.as_encoded())
22        let encoded = self.as_encoded();
23        if encoded <= 0xFF {
24            write!(f, "0x{:02X}", encoded)
25        } else if encoded <= 0xFFFF {
26            write!(f, "0x{:04X}", encoded)
27        } else if encoded <= 0xFFFFFF {
28            write!(f, "0x{:06X}", encoded)
29        } else if encoded <= 0xFFFFFFFF {
30            write!(f, "0x{:08X}", encoded)
31        } else if encoded <= 0xFFFFFFFFFF {
32            write!(f, "0x{:010X}", encoded)
33        } else if encoded <= 0xFFFFFFFFFFFF {
34            write!(f, "0x{:012X}", encoded)
35        } else if encoded <= 0xFFFFFFFFFFFFFF {
36            write!(f, "0x{:014X}", encoded)
37        } else {
38            write!(f, "0x{:016X}", encoded)
39        }
40    }
41}
42impl Debug for VInt64 {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        let mut t = f.debug_struct("VInt64");
45        if !self.is_unknown {
46            t.field("value", &self.value);
47        } else {
48            t.field("value", &"Unknown");
49        }
50        t.field("memory", &format!("{}", self));
51        t.finish()
52    }
53}
54
55impl Deref for VInt64 {
56    type Target = u64;
57    fn deref(&self) -> &Self::Target {
58        &self.value
59    }
60}
61impl VInt64 {
62    /// Create a VInt64 from an already encoded u64 value.
63    pub const fn from_encoded(enc: u64) -> Self {
64        if enc == 0xFF {
65            Self {
66                value: 127,
67                is_unknown: true,
68            }
69        } else if enc == 0x407F {
70            Self {
71                value: 127,
72                is_unknown: false,
73            }
74        } else {
75            Self {
76                value: enc & (u64::MAX >> (enc.leading_zeros() + 1)),
77                is_unknown: false,
78            }
79        }
80    }
81
82    /// Create a VInt64 representing an unknown size.
83    pub const fn new_unknown() -> Self {
84        Self {
85            value: 127,
86            is_unknown: true,
87        }
88    }
89
90    /// Create a VInt64 from a u64 value.
91    pub const fn new(value: u64) -> Self {
92        Self {
93            value,
94            is_unknown: false,
95        }
96    }
97
98    /// Create a VInt64 from an already encoded u64 value.
99    pub fn as_encoded(&self) -> u64 {
100        if self.is_unknown {
101            return 0xFF;
102        }
103        if self.value == 127 {
104            return 0x407F;
105        }
106
107        let size = VInt64::encode_size(self.value);
108        let mut sbuf = [0u8; 8];
109        let slice = &mut sbuf[8 - size..];
110        slice.copy_from_slice(&self.value.to_be_bytes()[8 - size..]);
111        slice[0] |= 1u8 << (8 - size);
112        u64::from_be_bytes(sbuf)
113    }
114
115    /// Get the size in bytes of the encoded representation of a u64 value.
116    pub const fn encode_size(value: u64) -> usize {
117        let leading_zeros = value.leading_zeros() as usize;
118        let total_bits = 64 - leading_zeros;
119        if total_bits == 0 {
120            1
121        } else {
122            (total_bits + 6).div_euclid(7)
123        }
124    }
125}
126
127impl ReadFrom for VInt64 {
128    fn read_from<R: std::io::Read + ?Sized>(r: &mut R) -> crate::Result<Self> {
129        let mut first_byte_buf = [0u8; 1];
130        r.read_exact(&mut first_byte_buf)?;
131        let first_byte = first_byte_buf[0];
132        if first_byte == 0xFF {
133            return Ok(VInt64 {
134                value: 127,
135                is_unknown: true,
136            });
137        }
138
139        let leading_zeros = first_byte.leading_zeros() as usize;
140        if leading_zeros >= 8 {
141            return Err(crate::error::Error::InvalidVInt);
142        }
143
144        if leading_zeros == 0 {
145            Ok(VInt64 {
146                value: (first_byte & 0b0111_1111) as u64,
147                is_unknown: false,
148            })
149        } else {
150            let mut buf = [0u8; 8];
151            let read_buf = &mut buf[8 - leading_zeros..];
152            r.read_exact(read_buf)?;
153            if leading_zeros != 7 {
154                buf[8 - leading_zeros - 1] = first_byte & (0xFF >> (leading_zeros + 1));
155            }
156            Ok(VInt64 {
157                value: u64::from_be_bytes(buf),
158                is_unknown: false,
159            })
160        }
161    }
162}
163
164#[cfg(feature = "tokio")]
165#[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
166impl crate::io::tokio_impl::AsyncReadFrom for VInt64 {
167    async fn async_read_from<R: tokio::io::AsyncRead + Unpin + ?Sized>(
168        r: &mut R,
169    ) -> crate::Result<Self> {
170        let mut first_byte_buf = [0u8; 1];
171        tokio::io::AsyncReadExt::read_exact(r, &mut first_byte_buf).await?;
172        let first_byte = first_byte_buf[0];
173        if first_byte == 0xFF {
174            return Ok(VInt64 {
175                value: 127,
176                is_unknown: true,
177            });
178        }
179
180        let leading_zeros = first_byte.leading_zeros() as usize;
181        if leading_zeros >= 8 {
182            return Err(crate::error::Error::InvalidVInt);
183        }
184
185        if leading_zeros == 0 {
186            Ok(VInt64 {
187                value: (first_byte & 0b0111_1111) as u64,
188                is_unknown: false,
189            })
190        } else {
191            let mut buf = [0u8; 8];
192            let read_buf = &mut buf[8 - leading_zeros..];
193            tokio::io::AsyncReadExt::read_exact(r, read_buf).await?;
194            if leading_zeros != 7 {
195                buf[8 - leading_zeros - 1] = first_byte & (0xFF >> (leading_zeros + 1));
196            }
197            Ok(VInt64 {
198                value: u64::from_be_bytes(buf),
199                is_unknown: false,
200            })
201        }
202    }
203}
204
205impl Decode for VInt64 {
206    fn decode<B: Buf>(buf: &mut B) -> crate::Result<Self> {
207        let first_byte = buf.try_get_u8()?;
208        if first_byte == 0 {
209            return Err(Error::InvalidVInt);
210        }
211        if first_byte == 0xFF {
212            return Ok(VInt64 {
213                value: 127,
214                is_unknown: true,
215            });
216        }
217        let leading_zeros = first_byte.leading_zeros() as usize;
218
219        if leading_zeros == 0 {
220            Ok(VInt64 {
221                value: (first_byte & 0b0111_1111) as u64,
222                is_unknown: false,
223            })
224        } else {
225            if buf.remaining() < leading_zeros {
226                return Err(Error::try_get_error(leading_zeros, buf.remaining()));
227            }
228            let mut bytes = [0u8; 8];
229            let read_buf = &mut bytes[8 - leading_zeros..];
230            read_buf.copy_from_slice(&buf.chunk()[..leading_zeros]);
231
232            if leading_zeros != 7 {
233                bytes[8 - leading_zeros - 1] = first_byte & (0xFF >> (leading_zeros + 1));
234            }
235            buf.advance(leading_zeros);
236            Ok(VInt64 {
237                value: u64::from_be_bytes(bytes),
238                is_unknown: false,
239            })
240        }
241    }
242}
243
244impl Encode for VInt64 {
245    fn encode<B: BufMut>(&self, buf: &mut B) -> crate::Result<()> {
246        if self.is_unknown {
247            buf.put_slice(&[0xFF]);
248            return Ok(());
249        }
250        if self.value == 127 {
251            buf.put_slice(&[0x40, 0x7F]);
252            return Ok(());
253        }
254
255        let size = VInt64::encode_size(self.value);
256        let mut sbuf = [0u8; 8];
257        let slice = &mut sbuf[8 - size..];
258        slice.copy_from_slice(&self.value.to_be_bytes()[8 - size..]);
259        slice[0] |= 1u8 << (8 - size);
260        buf.put_slice(slice);
261        Ok(())
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268    use std::convert::TryInto;
269
270    #[test]
271    fn test_encode_size() {
272        let test_pair = [
273            (vec![0b1000_0000], 0),
274            (vec![0b1000_0001], 1),
275            (vec![0b0100_0000, 0xFF], 0xFF),
276            (vec![0b0100_0001, 0xFF], 0b1_1111_1111),
277            (vec![0b0111_1111, 0xFF], 0b11_1111_1111_1111),
278            (vec![0b0010_0000, 0b0111_1111, 0xFF], 0b111_1111_1111_1111),
279            (vec![0b0010_0000, 0xFF, 0xFF], 0xFFFF),
280            (vec![0b0011_1111, 0xFF, 0xFF], 0b1_1111_1111_1111_1111_1111),
281            (
282                vec![1, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF],
283                0xFF_FFFF_FFFF_FFFF,
284            ),
285        ];
286        for (encoded, val) in test_pair {
287            assert_eq!(VInt64::encode_size(val), encoded.len());
288        }
289    }
290
291    #[test]
292    fn test_encode() {
293        let test_pair = [
294            (vec![0b1000_0000], 0),
295            (vec![0b1000_0001], 1),
296            (vec![0b0100_0000, 0xFF], 0xFF),
297            (vec![0b0100_0001, 0xFF], 0b1_1111_1111),
298            (vec![0b0111_1111, 0xFF], 0b11_1111_1111_1111),
299            (vec![0b0010_0000, 0b0111_1111, 0xFF], 0b111_1111_1111_1111),
300            (vec![0b0010_0000, 0xFF, 0xFF], 0xFFFF),
301            (vec![0b0011_1111, 0xFF, 0xFF], 0b1_1111_1111_1111_1111_1111),
302            (
303                vec![1, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF],
304                0xFF_FFFF_FFFF_FFFF,
305            ),
306        ];
307        for (encoded, val) in test_pair {
308            let v = VInt64 {
309                value: val,
310                is_unknown: false,
311            };
312            let mut out = vec![];
313            v.encode(&mut out).unwrap();
314            assert_eq!(encoded, out);
315
316            let encoded_num = v.as_encoded();
317            let mut enc8 = vec![0u8; 8 - encoded.len()];
318            enc8.extend_from_slice(&encoded);
319            let encoded_from = u64::from_be_bytes(enc8.try_into().unwrap());
320            assert_eq!(encoded_num, encoded_from);
321        }
322    }
323
324    #[test]
325    fn test_decode() {
326        let test_pair = [
327            (vec![0b1000_0000], 0),
328            (vec![0b1000_0001], 1),
329            (vec![0b0100_0000, 0xFF], 0xFF),
330            (vec![0b0100_0001, 0xFF], 0b1_1111_1111),
331            (vec![0b0111_1111, 0xFF], 0b11_1111_1111_1111),
332            (vec![0b0010_0000, 0b0111_1111, 0xFF], 0b111_1111_1111_1111),
333            (vec![0b0010_0000, 0xFF, 0xFF], 0xFFFF),
334            (vec![0b0011_1111, 0xFF, 0xFF], 0b1_1111_1111_1111_1111_1111),
335            (
336                vec![1, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF],
337                0xFF_FFFF_FFFF_FFFF,
338            ),
339        ];
340        for (encoded, val) in test_pair {
341            // test read
342            let mut c = std::io::Cursor::new(encoded.clone());
343            let vint = VInt64::read_from(&mut c).unwrap();
344            assert_eq!(*vint, val);
345
346            // test decode
347            let ecoded2 = encoded.clone();
348            let mut slice_encoded2 = Bytes::from(ecoded2);
349            let vint_decoded = VInt64::decode(&mut slice_encoded2).unwrap();
350            assert_eq!(*vint_decoded, val);
351
352            // test from_encoded
353            let mut enc8 = vec![0u8; 8 - encoded.len()];
354            enc8.extend_from_slice(&encoded);
355            let v = VInt64::from_encoded(u64::from_be_bytes(enc8.try_into().unwrap()));
356            assert_eq!(*v, val);
357        }
358    }
359
360    #[test]
361    fn test_unknown() {
362        let v1 = VInt64::read_from(&mut std::io::Cursor::new(vec![0xFF])).unwrap();
363        let vv1 = VInt64::from_encoded(0xFF);
364        assert!(v1.is_unknown);
365        assert!(vv1.is_unknown);
366
367        let v2 = VInt64::read_from(&mut std::io::Cursor::new(vec![0x80])).unwrap();
368        let vv2 = VInt64::from_encoded(0x80);
369
370        assert!(!v2.is_unknown);
371        assert!(!vv2.is_unknown);
372
373        let v3 = VInt64::read_from(&mut std::io::Cursor::new(vec![0x40, 0x7F])).unwrap();
374        let vv3 = VInt64::from_encoded(0x407F);
375        assert_eq!(*v3, 127);
376        assert_eq!(*vv3, 127);
377
378        assert_ne!(VInt64::new(127), VInt64::new_unknown());
379        assert_eq!(VInt64::new(127).as_encoded(), 0x407F);
380    }
381}
382
383/// EBML element header, consisting of an ID and a size.
384#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
385pub struct Header {
386    /// EBML ID of the element.
387    pub id: VInt64,
388    /// Size of the element's data, excluding the header itself.
389    pub size: VInt64,
390}
391
392impl ReadFrom for Header {
393    fn read_from<R: std::io::Read + ?Sized>(reader: &mut R) -> crate::Result<Self> {
394        let id = VInt64::read_from(reader)?;
395        let size = VInt64::read_from(reader)?;
396        Ok(Self { id, size })
397    }
398}
399
400#[cfg(feature = "tokio")]
401#[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
402impl crate::io::tokio_impl::AsyncReadFrom for Header {
403    async fn async_read_from<R: tokio::io::AsyncRead + Unpin + ?Sized>(
404        r: &mut R,
405    ) -> crate::Result<Self> {
406        let id = VInt64::async_read_from(r).await?;
407        let size = VInt64::async_read_from(r).await?;
408        Ok(Self { id, size })
409    }
410}
411
412impl Decode for Header {
413    fn decode<B: Buf>(buf: &mut B) -> crate::Result<Self> {
414        let id = VInt64::decode(buf)?;
415        let size = VInt64::decode(buf)?;
416        Ok(Self { id, size })
417    }
418}
419
420impl Encode for Header {
421    fn encode<B: BufMut>(&self, buf: &mut B) -> crate::Result<()> {
422        self.id.encode(buf)?;
423        self.size.encode(buf)?;
424        Ok(())
425    }
426}