Skip to main content

mkv_element/
base.rs

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